Skip to content

fix(rocm): resolve segfault during inference on ROCm/HIP backends (#163) - #1929

Draft
ljluestc wants to merge 1 commit into
haotian-liu:mainfrom
ljluestc:private/issue-163-rocm-segfault-fix
Draft

fix(rocm): resolve segfault during inference on ROCm/HIP backends (#163)#1929
ljluestc wants to merge 1 commit into
haotian-liu:mainfrom
ljluestc:private/issue-163-rocm-segfault-fix

Conversation

@ljluestc

Copy link
Copy Markdown

Three interacting bugs caused a SIGSEGV when running model inference on AMD GPUs via PyTorch ROCm (HIP):

  1. torch.inference_mode() is thread-local. generate_stream() is decorated with @torch.inference_mode() but model.generate() is dispatched onto a child Thread whose HIP context is independent. Without inference mode active, PyTorch/HIP attempts to allocate gradient-tracking tensors inside the generation kernel, corrupting the HIP memory pool and triggering a segmentation fault. Fixed: _generate_with_inference_mode() closure wraps model.generate inside an explicit torch.inference_mode() context.

  2. The HIP memory allocator uses an aggressive split-and-reuse policy by default. For large language model inference this leads to excessive fragmentation on ROCm 5.x. Fixed: configure_hip_allocator() sets PYTORCH_HIP_ALLOC_CONF via os.environ.setdefault at startup (user-overridable).

  3. The device string 'rocm' is not recognised by PyTorch, which expects 'cuda' even on ROCm builds. Passing --device rocm silently fell through to incorrect device-map handling. Fixed: 'rocm' is normalised to 'cuda' in both model_worker.py and builder.py.

New files

llava/serve/rocm_utils.py — lightweight is_rocm() / configure_hip_allocator()
helpers with no heavy model dependencies
tests/test_rocm_segfault_fix.py — 22 unit tests (all pass)

…otian-liu#163)

Three interacting bugs caused a SIGSEGV when running model inference on
AMD GPUs via PyTorch ROCm (HIP):

1. torch.inference_mode() is thread-local.  generate_stream() is
   decorated with @torch.inference_mode() but model.generate() is
   dispatched onto a child Thread whose HIP context is independent.
   Without inference mode active, PyTorch/HIP attempts to allocate
   gradient-tracking tensors inside the generation kernel, corrupting
   the HIP memory pool and triggering a segmentation fault.
   Fixed: _generate_with_inference_mode() closure wraps model.generate
   inside an explicit torch.inference_mode() context.

2. The HIP memory allocator uses an aggressive split-and-reuse policy
   by default.  For large language model inference this leads to
   excessive fragmentation on ROCm 5.x.
   Fixed: configure_hip_allocator() sets PYTORCH_HIP_ALLOC_CONF via
   os.environ.setdefault at startup (user-overridable).

3. The device string 'rocm' is not recognised by PyTorch, which expects
   'cuda' even on ROCm builds.  Passing --device rocm silently fell
   through to incorrect device-map handling.
   Fixed: 'rocm' is normalised to 'cuda' in both model_worker.py and
   builder.py.

New files
---------
llava/serve/rocm_utils.py   — lightweight is_rocm() / configure_hip_allocator()
                              helpers with no heavy model dependencies
tests/test_rocm_segfault_fix.py — 22 unit tests (all pass)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant