Repository navigation
Profiling Torch.compile (with_stack=True) leads to frozen program with massive memory. #136404
Description
🐛 Describe the bug
I am using gemma-2b with vllm. If I torch.compile(fullgraph=False, backend='inductor') and then profile model.generate in vllm using:
with torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
with_stack=True,
experimental_config=torch._C._profiler._ExperimentalConfig(verbose=True),
) as prof:
the profiler runs forever and shows huge memory load. If I profile with
with torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
with_modules=True,
experimental_config=torch._C._profiler._ExperimentalConfig(verbose=True),
) as prof:
the job finishes in ~40 seconds.
Error logs
No response
Minified repro
No response
Versions
torch==2.4.0+cu118
torchmetrics>=1.4.2
vllm @ https://github.com/vllm-project/vllm/releases/download/v0.5.4/vllm-0.5.4+cu118-cp310-cp310-manylinux1_x86_64.whl
vllm-flash-attn @ https://github.com/vllm-project/flash-attention/releases/download/v2.6.1/vllm_flash_attn-2.6.1+cu118-cp310-cp310-manylinux1_x86_64.whl
cc @robieta @chaekit @guotuofeng @guyang3532 @dzhulgakov @davidberard98 @briancoutinho @sraikund16 @sanrise @chauhang @penguinwu @zou3519 @aaronenyeshi @ezyang
Activity
Etragas commentedon Sep 22, 2024
Example of the memory load increasing from the profiler:
10 minutes later
Eventually, this process will consume the 1 TB of ram available on my machine and exit with no message.
[-]Torch.compile + torch.profiler (with_stack=True) leads to frozen program with massive memory.[/-][+]Profiling Torch.compile (with_stack=True) leads to frozen program with massive memory.[/+]Etragas commentedon Sep 22, 2024
I have diagnosed the core issue. I was calling model.generate for the first time insde my profiling block. When this happens we trigger the lazy compilation of dynamo+inductor. This causes the pytorch profiler to try and profile the compilation, and this doesn't go well for the profiler...
I'm not sure if it's supported to use pytorch profiler to profile the actual compilation task.
If not, is there a way to add a warning / error to the user to force compilation by running the task first, and then profile the compiled code?
ezyang commentedon Sep 23, 2024
disabling profiler when we hit compile regions is very reasonable
Etragas commentedon Sep 23, 2024
@ezyang do you have any suggestions for a scrappy implementation of this? We are trying to profile a model that does recompilation on each forward pass so it's tricky to force up front compilation.
ezyang commentedon Sep 24, 2024
if you know the context manager to disable profiling slap it in
in torch/_dynamo/convert_frame.py, maybe anywhere else compile_lock is taken too
Disable profiler Python tracing during compile