Skip to content

Set PYTORCH_CUDA_ALLOC_CONF before platform init - #1595

Open
Shan2L wants to merge 1 commit into
ModelTC:mainfrom
Shan2L:fix/rocm-alloc-conf-before-platform-init
Open

Shan2L wants to merge 1 commit into
ModelTC:mainfrom
Shan2L:fix/rocm-alloc-conf-before-platform-init

Conversation

@Shan2L

@Shan2L Shan2L commented Oct 10, 2026

Copy link
Copy Markdown

Fixes #1593

lightx2v/pipeline.py sets PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True via setdefault, but
lightx2v/__init__.py imports lightx2v_platform.set_ai_device before pipeline.py. Some platform
backends initialize the CUDA/HIP caching allocator during that import (on amd_rocm, importing aiter
does), and after that the allocator no longer reads the variable. As a result amd_rocm runs without
expandable segments, and MiniMax-H3 OOMs on 32 GB GPUs.

This PR applies the same setdefault at the top of lightx2v/__init__.py, before the platform import.
The existing line in pipeline.py is unchanged. An explicitly exported PYTORCH_CUDA_ALLOC_CONF still
wins. Platforms whose allocator does not read this variable are unaffected.

Test

  • is_expandable of the first segment after import lightx2v: amd_rocm False -> True; cuda True (unchanged).
  • MiniMax-H3 ref2av TP2 (768x1344, 362 frames, 3 steps, block offload) on 2x R9700, PLATFORM=amd_rocm, main:
    OOM at step 1 -> completes; 4 runs on 4 GPU pairs, 100.0-101.8 s/step (same as PLATFORM=cuda).
  • ruff check / ruff format --check (v0.11.0) clean.

lightx2v/__init__.py imports lightx2v_platform.set_ai_device before
lightx2v/pipeline.py sets PYTORCH_CUDA_ALLOC_CONF. On amd_rocm the platform
import pulls in aiter, which initializes the caching allocator, so
expandable_segments was silently ignored and MiniMax-H3 OOMed on 32 GB GPUs.

Signed-off-by: Sylvan Liu <Sylvan.Liu@amd.com>
@Shan2L

Shan2L commented Oct 11, 2026

Copy link
Copy Markdown
Author

@helloyongyang @gushiqiao @llmc-reviewer could you take a look, or point me to the right reviewer? This is a 6-line fix for #1593: on amd_rocm the allocator is initialized before pipeline.py sets PYTORCH_CUDA_ALLOC_CONF, so MiniMax-H3 OOMs on 32 GB GPUs. CUDA behavior is unchanged and an exported value still wins.

It is one of two ROCm prerequisites (with #1596) for the gfx1201 attention backends in #1598 / #1599 (integration plan: ROCm/aiter#6292). Happy to adjust anything.

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.

[Bug] PLATFORM=amd_rocm ignores expandable_segments:True, MiniMax-H3 OOMs on 32 GB GPUs

1 participant