Fix by avoiding failure for triton kernels in pickling error

[nemanjaudovic]

Motivation

The following error was encountered when running some models using pytorch 2.12 and vllm:

_pickle.PicklingError: Can't pickle <function launcher at 0x7fbd1416e480>:
attribute lookup launcher on __main__ failed

It is present on release/2.12 branch.

Technical Details

Fix for this error already exists on upstream pytorch repo.
Cherry-pick of a5791f3 from upstream pytorch.

Test Plan

Error was encountered running this command:

 vllm bench throughput \
 --model mistralai/Mistral-7B-Instruct-v0.2 \
 --max_model_len 4096 \
 --trust-remote-code \
 --dataset-name sharegpt \
 --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
 --num_prompts 1000 \
 --tensor_parallel_size 1  \
 --gpu-memory-utilization 0.8

Test Result

After applying this cherry pick workload runs without crashing, and AOT cache is able to be saved.

[rocm-repo-management-api]

Jenkins build for 9023ac176f4814cbfb1c79025bd254f54291c7df commit finished as FAILURE
Links: Pipeline Overview / Build artifacts / Test Results

Detected error during base docker image building:

#38 17.98 Cloning into 'drm'...
#38 21.00 /drm /
#38 21.00 patching file amdgpu/amdgpu_asic_id.c
#38 21.00 ./install_rocm_drm.sh: line 154: meson: command not found
#38 21.00 ./install_rocm_drm.sh: line 155: pushd: builddir: No such file or directory
#38 21.05 ninja: error: loading 'build.ninja': No such file or directory
#38 21.05 /
#38 21.05 ./install_rocm_drm.sh: line 159: popd: directory stack empty
#38 ERROR: process "/bin/sh -c if [ \"${ROCM_VERSION}\" != \"nightly\" ]; then bash ./install_rocm_drm.sh /usr ; fi && rm install_rocm_drm.sh" did not complete successfully: exit code: 1
------
 > [33/54] RUN if [ "7.2.3" != "nightly" ]; then bash ./install_rocm_drm.sh /usr ; fi && rm install_rocm_drm.sh: