throughputbenchmark

[ghstack-poisoned]

…hread in"

TorchDynamo will guard grad_mode and the local dispatch key set.
https://github.com/pytorch/pytorch/blob/3a429423fcf72430e7a36c79e263c877d7a4ef72/torch/csrc/dynamo/guards.cpp#L13-L16

While using ThroughputBenchmark, those tls state will not be init as same as the main thread status.
https://github.com/pytorch/pytorch/blob/3a429423fcf72430e7a36c79e263c877d7a4ef72/torch/csrc/utils/throughput_benchmark-inl.h#L64-L94

Run following scripts
```
import torch
linear = torch.nn.Linear(128, 128)
compiled = torch.compile(linear)
x = torch.rand(10, 128)
with torch.no_grad(), torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    compiled(x)
    compiled(x)

from torch._dynamo import config
config.error_on_recompile = True
from torch.utils import ThroughputBenchmark
with torch.no_grad(), torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    bench = ThroughputBenchmark(compiled)
    bench.add_input(x)
    stats = bench.benchmark(
        num_calling_threads=10,
        num_warmup_iters=100,
        num_iters=100,
    )
    print(stats)
```
will lead to 2 re-compile reasons:
```
triggered by the following guard failure(s): ___check_global_state()
triggered by the following guard failure(s): tensor 'x' dispatch key set mismatch.
```

This will trigger a re-compile in torchdynamo. But since `ThroughputBenchmark` is used for sharing weight within threads, the model should not be changed anymore while running the benchmark. So this PR is to init the tls state as same as main thread. Then we can use ` ThroughputBenchmark` to run torchdynamo optimized models.




throughputbenchmark

[ghstack-poisoned]

…hread in"

TorchDynamo will guard grad_mode and the local dispatch key set.
https://github.com/pytorch/pytorch/blob/3a429423fcf72430e7a36c79e263c877d7a4ef72/torch/csrc/dynamo/guards.cpp#L13-L16

While using ThroughputBenchmark, those tls state will not be init as same as the main thread status.
https://github.com/pytorch/pytorch/blob/3a429423fcf72430e7a36c79e263c877d7a4ef72/torch/csrc/utils/throughput_benchmark-inl.h#L64-L94

Run following scripts
```
import torch
linear = torch.nn.Linear(128, 128)
compiled = torch.compile(linear)
x = torch.rand(10, 128)
with torch.no_grad(), torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    compiled(x)
    compiled(x)

from torch._dynamo import config
config.error_on_recompile = True
from torch.utils import ThroughputBenchmark
with torch.no_grad(), torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    bench = ThroughputBenchmark(compiled)
    bench.add_input(x)
    stats = bench.benchmark(
        num_calling_threads=10,
        num_warmup_iters=100,
        num_iters=100,
    )
    print(stats)
```
will lead to 2 re-compile reasons:
```
triggered by the following guard failure(s): ___check_global_state()
triggered by the following guard failure(s): tensor 'x' dispatch key set mismatch.
```

This will trigger a re-compile in torchdynamo. But since `ThroughputBenchmark` is used for sharing weight within threads, the model should not be changed anymore while running the benchmark. So this PR is to init the tls state as same as main thread. Then we can use ` ThroughputBenchmark` to run torchdynamo optimized models.




throughputbenchmark

[ghstack-poisoned]

…hread in"

TorchDynamo will guard grad_mode and the local dispatch key set.
https://github.com/pytorch/pytorch/blob/3a429423fcf72430e7a36c79e263c877d7a4ef72/torch/csrc/dynamo/guards.cpp#L13-L16

While using ThroughputBenchmark, those tls state will not be init as same as the main thread status.
https://github.com/pytorch/pytorch/blob/3a429423fcf72430e7a36c79e263c877d7a4ef72/torch/csrc/utils/throughput_benchmark-inl.h#L64-L94

Run following scripts
```
import torch
linear = torch.nn.Linear(128, 128)
compiled = torch.compile(linear)
x = torch.rand(10, 128)
with torch.no_grad(), torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    compiled(x)
    compiled(x)

from torch._dynamo import config
config.error_on_recompile = True
from torch.utils import ThroughputBenchmark
with torch.no_grad(), torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    bench = ThroughputBenchmark(compiled)
    bench.add_input(x)
    stats = bench.benchmark(
        num_calling_threads=10,
        num_warmup_iters=100,
        num_iters=100,
    )
    print(stats)
```
will lead to 2 re-compile reasons:
```
triggered by the following guard failure(s): ___check_global_state()
triggered by the following guard failure(s): tensor 'x' dispatch key set mismatch.
```

This will trigger a re-compile in torchdynamo. But since `ThroughputBenchmark` is used for sharing weight within threads, the model should not be changed anymore while running the benchmark. So this PR is to init the tls state as same as main thread. Then we can use ` ThroughputBenchmark` to run torchdynamo optimized models.




throughputbenchmark

[ghstack-poisoned]

…hread in"

TorchDynamo will guard grad_mode and the local dispatch key set.
https://github.com/pytorch/pytorch/blob/3a429423fcf72430e7a36c79e263c877d7a4ef72/torch/csrc/dynamo/guards.cpp#L13-L16

While using ThroughputBenchmark, those tls state will not be init as same as the main thread status.
https://github.com/pytorch/pytorch/blob/3a429423fcf72430e7a36c79e263c877d7a4ef72/torch/csrc/utils/throughput_benchmark-inl.h#L64-L94

Run following scripts
```
import torch
linear = torch.nn.Linear(128, 128)
compiled = torch.compile(linear)
x = torch.rand(10, 128)
with torch.no_grad(), torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    compiled(x)
    compiled(x)

from torch._dynamo import config
config.error_on_recompile = True
from torch.utils import ThroughputBenchmark
with torch.no_grad(), torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    bench = ThroughputBenchmark(compiled)
    bench.add_input(x)
    stats = bench.benchmark(
        num_calling_threads=10,
        num_warmup_iters=100,
        num_iters=100,
    )
    print(stats)
```
will lead to 2 re-compile reasons:
```
triggered by the following guard failure(s): ___check_global_state()
triggered by the following guard failure(s): tensor 'x' dispatch key set mismatch.
```

This will trigger a re-compile in torchdynamo. But since `ThroughputBenchmark` is used for sharing weight within threads, the model should not be changed anymore while running the benchmark. So this PR is to init the tls state as same as main thread. Then we can use ` ThroughputBenchmark` to run torchdynamo optimized models.




throughputbenchmark

[ghstack-poisoned]