Repository: pytorch
Issue: inductor error when torch.compile on distrifuser
Labels: triaged,oncall: pt2,module: inductor

### 🐛 Describe the bug

I applied torch.compile on [distrifuser](https://github.com/mit-han-lab/distrifuser/tree/main). In particular, I patched this [code](https://github.com/mit-han-lab/distrifuser/blob/main/distrifuser/pipelines.py) by applying torch.compile on its [unet](https://www.internalfb.com/phabricator/paste/view/P1397326708).

`torch.compile(unet, backend="eager")` and `torch.compile(unet, backend="aot_eager")` work. However, `torch.compile(unet, backend="inductor")` gives an error `LLVM ERROR: Failed to translate TritonGPU to LLVM IR`.






### Error logs

(distrifuser) [boyuan@devvm4377.nao0 ~/local/distrifuser (main)]$ torchrun --nproc_per_node=1 scripts/run_sdxl.py --mode benchmark --output_type latent
/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/diffusers/utils/outputs.py:63: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.
  torch.utils._pytree._register_pytree_node(
/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(
/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/diffusers/utils/outputs.py:63: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.
  torch.utils._pytree._register_pytree_node(
Cannot initialize model with low cpu memory usage because `accelerate` was not found in the environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install `accelerate` for faster and less memory-intense model loading. You can do so with: 
```
pip install accelerate
```
.
/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
  warnings.warn(
Cannot initialize model with low cpu memory usage because `accelerate` was not found in the environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install `accelerate` for faster and less memory-intense model loading. You can do so with: 
```
pip install accelerate
```
.
Loading pipeline components...: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:09<00:00,  1.42s/it]
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
Failed to load Translate to LLVM IR failedLLVM ERROR: Failed to translate TritonGPU to LLVM IR.
[rank0]: Traceback (most recent call last):
[rank0]:   File "/data/users/boyuan/distrifuser/scripts/run_sdxl.py", line 159, in <module>
[rank0]:     main()
[rank0]:   File "/data/users/boyuan/distrifuser/scripts/run_sdxl.py", line 105, in main
[rank0]:     pipeline = DistriSDXLPipeline.from_pretrained(
[rank0]:   File "/data/users/boyuan/distrifuser/distrifuser/pipelines.py", line 46, in from_pretrained
[rank0]:     return DistriSDXLPipeline(pipeline, distri_config)
[rank0]:   File "/data/users/boyuan/distrifuser/distrifuser/pipelines.py", line 17, in __init__
[rank0]:     self.prepare()
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/data/users/boyuan/distrifuser/distrifuser/pipelines.py", line 149, in prepare
[rank0]:     pipeline.unet(**static_inputs, return_dict=False, record=True)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/data/users/boyuan/distrifuser/distrifuser/models/distri_sdxl_unet_pp.py", line 119, in forward
[rank0]:     output = self.model(
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 451, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 921, in catch_errors
[rank0]:     return callback(frame, cache_entry, hooks, frame_state, skip=1)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 786, in _convert_frame
[rank0]:     result = inner_convert(
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert
[rank0]:     return _compile(
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/contextlib.py", line 79, in inner
[rank0]:     return func(*args, **kwds)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 676, in _compile
[rank0]:     guarded_code = compile_inner(code, one_graph, hooks, transform)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
[rank0]:     r = func(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 535, in compile_inner
[rank0]:     out_code = transform_code_object(code, transform)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1036, in transform_code_object
[rank0]:     transformations(instructions, code_options)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 165, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 500, in transform
[rank0]:     tracer.run()
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2149, in run
[rank0]:     super().run()
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run
[rank0]:     and self.step()
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step
[rank0]:     getattr(self, inst.opname)(inst)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2268, in RETURN_VALUE
[rank0]:     self.output.compile_subgraph(
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1001, in compile_subgraph
[rank0]:     self.compile_and_call_fx_graph(tx, pass2.graph_output_vars(), root)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/contextlib.py", line 79, in inner
[rank0]:     return func(*args, **kwds)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1178, in compile_and_call_fx_graph
[rank0]:     compiled_fn = self.call_user_compiler(gm)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
[rank0]:     r = func(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1251, in call_user_compiler
[rank0]:     raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1232, in call_user_compiler
[rank0]:     compiled_fn = compiler_fn(gm, self.example_inputs())
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/repro/after_dynamo.py", line 117, in debug_wrapper
[rank0]:     compiled_gm = compiler_fn(gm, example_inputs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/__init__.py", line 1731, in __call__
[rank0]:     return compile_fx(model_, inputs_, config_patches=self.config)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/contextlib.py", line 79, in inner
[rank0]:     return func(*args, **kwds)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1330, in compile_fx
[rank0]:     return aot_autograd(
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/backends/common.py", line 58, in compiler_fn
[rank0]:     cg = aot_module_simplified(gm, example_inputs, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 903, in aot_module_simplified
[rank0]:     compiled_fn = create_aot_dispatcher_function(
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
[rank0]:     r = func(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 628, in create_aot_dispatcher_function
[rank0]:     compiled_fn = compiler_fn(flat_fn, fake_flat_args, aot_config, fw_metadata=fw_metadata)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 443, in aot_wrapper_dedupe
[rank0]:     return compiler_fn(flat_fn, leaf_flat_args, aot_config, fw_metadata=fw_metadata)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 648, in aot_wrapper_synthetic_base
[rank0]:     return compiler_fn(flat_fn, flat_args, aot_config, fw_metadata=fw_metadata)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 119, in aot_dispatch_base
[rank0]:     compiled_fw = compiler(fw_module, updated_flat_args)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
[rank0]:     r = func(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 1257, in fw_compiler_base
[rank0]:     return inner_compile(
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/repro/after_aot.py", line 83, in debug_wrapper
[rank0]:     inner_compiled_fn = compiler_fn(gm, example_inputs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_inductor/debug.py", line 304, in inner
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/contextlib.py", line 79, in inner
[rank0]:     return func(*args, **kwds)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/contextlib.py", line 79, in inner
[rank0]:     return func(*args, **kwds)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
[rank0]:     r = func(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 438, in compile_fx_inner
[rank0]:     compiled_graph = fx_codegen_and_compile(
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 714, in fx_codegen_and_compile
[rank0]:     compiled_fn = graph.compile_to_fn()
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_inductor/graph.py", line 1307, in compile_to_fn
[rank0]:     return self.compile_to_module().call
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
[rank0]:     r = func(*args, **kwargs)
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_inductor/graph.py", line 1254, in compile_to_module
[rank0]:     mod = PyCodeCache.load_by_key_path(
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_inductor/codecache.py", line 2160, in load_by_key_path
[rank0]:     exec(code, mod.__dict__, mod.__dict__)
[rank0]:   File "/tmp/torchinductor_boyuan/xz/cxzhj5q4vgrlnstqsb3f6jvg2lloaxi2vbxclitjx5jmyoayneii.py", line 6489, in <module>
[rank0]:     async_compile.wait(globals())
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_inductor/codecache.py", line 2715, in wait
[rank0]:     scope[key] = result.result()
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/site-packages/torch/_inductor/codecache.py", line 2522, in result
[rank0]:     self.future.result()
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/concurrent/futures/_base.py", line 458, in result
[rank0]:     return self.__get_result()
[rank0]:   File "/home/boyuan/.conda/envs/distrifuser/lib/python3.10/concurrent/futures/_base.py", line 403, in __get_result
[rank0]:     raise self._exception
[rank0]: torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
[rank0]: BrokenProcessPool: A process in the process pool was terminated abruptly while the future was running or pending.

[rank0]: Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


[rank0]: You can suppress this exception and fall back to eager by setting:
[rank0]:     import torch._dynamo
[rank0]:     torch._dynamo.config.suppress_errors = True


### Minified repro

torchrun --nproc_per_node=4 scripts/run_sdxl.py --mode benchmark --output_type latent

### Versions

PyTorch version: 2.3.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: CentOS Stream 9 (x86_64)
GCC version: (GCC) 11.4.1 20231218 (Red Hat 11.4.1-3)
Clang version: Could not collect
CMake version: version 3.26.4
Libc version: glibc-2.34

Python version: 3.10.14 (main, May  6 2024, 19:42:50) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.19.0-0_fbk12_hardened_11583_g0bef9520ca2b-x86_64-with-glibc2.34
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA H100
GPU 1: NVIDIA H100
GPU 2: NVIDIA H100
GPU 3: NVIDIA H100

Nvidia driver version: 525.105.17
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Address sizes:                   52 bits physical, 57 bits virtual
Byte Order:                      Little Endian
CPU(s):                          184
On-line CPU(s) list:             0-183
Vendor ID:                       AuthenticAMD
Model name:                      AMD EPYC 9654 96-Core Processor
CPU family:                      25
Model:                           17
Thread(s) per core:              1
Core(s) per socket:              184
Socket(s):                       1
Stepping:                        1
BogoMIPS:                        4792.79
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm rep_good nopl cpuid extd_apicid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw perfctr_core invpcid_single ssbd ibrs ibpb stibp vmmcall fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx512_bf16 clzero xsaveerptr wbnoinvd arat npt lbrv nrip_save tsc_scale vmcb_clean pausefilter pfthreshold v_vmsave_vmload vgif avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid fsrm arch_capabilities
Virtualization:                  AMD-V
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       11.5 MiB (184 instances)
L1i cache:                       11.5 MiB (184 instances)
L2 cache:                        92 MiB (184 instances)
L3 cache:                        2.9 GiB (184 instances)
NUMA node(s):                    1
NUMA node0 CPU(s):               0-183
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP disabled, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] mypy==1.8.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] optree==0.11.0
[pip3] pytorch_sphinx_theme==0.0.24
[pip3] pytorch-triton==3.0.0+45fff310c8
[pip3] torch==2.3.1
[pip3] torchmetrics==1.4.0.post0
[pip3] torchprofile==0.0.4
[pip3] torchvision==0.18.1
[pip3] triton==2.3.1
[conda] blas                      1.0                         mkl  
[conda] mkl                       2023.1.0         h213fc3f_46344  
[conda] mkl-include               2023.1.0         h06a4308_46344  
[conda] mkl-service               2.4.0           py310h5eee18b_1  
[conda] mkl_fft                   1.3.8           py310h5eee18b_0  
[conda] mkl_random                1.2.4           py310hdb19cb5_0  
[conda] numpy                     1.26.4          py310h5f9d8c6_0  
[conda] numpy-base                1.26.4          py310hb5e798b_0  
[conda] optree                    0.11.0                   pypi_0    pypi
[conda] pytorch-sphinx-theme      0.0.24                    dev_0    <develop>
[conda] pytorch-triton            3.0.0+45fff310c8          pypi_0    pypi
[conda] torch                     2.3.1                    pypi_0    pypi
[conda] torchmetrics              1.4.0.post0              pypi_0    pypi
[conda] torchprofile              0.0.4                    pypi_0    pypi
[conda] torchvision               0.18.1                   pypi_0    pypi
[conda] triton                    2.3.1                    pypi_0    pypi

cc @ezyang @bdhirsh @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @msaroufim