feat: dynamic memory-aware GPU scheduling for model-free PTQ by rohan9446 · Pull Request #2976 · vllm-project/llm-compressor

rohan9446 · 2026-07-28T15:47:23Z

SUMMARY:

Fixes #2975, replaces the static round-robin GPU assignment in model-free PTQ with a dynamic, memory-aware scheduler.

Problem

PR #2773 introduced multi-GPU support by statically assigning shards to GPUs via devices[i % len(devices)]. When shards differ in size or GPUs differ in speed, the slower GPU accumulates pending work and eventually OOMs. This was observed on 6-GPU setups.

Solution

New scheduler.py module with a capacity-first scheduling loop:

Dynamic device selection: before submitting each job, the main thread calls torch.get_device_module(dev).mem_get_info() and picks the GPU with the most free VRAM
Reservation tracking: submitted-but-not-yet-allocated memory is tracked per device, preventing multiple jobs from stacking on the same GPU in a single scheduling pass
Forced fallback: if no device can fit any pending job and nothing is in flight, the smallest job is forced onto the most-free device with a warning (instead of silently dropping work)
Caching allocator flush: empty_cache() is called after harvesting each completed future so mem_get_info reflects actually-free memory
concurrent.futures.wait(FIRST_COMPLETED) replaces a hand-rolled threading.Event callback, eliminating a lost-wakeup race

_build_jobs() no longer takes a devices parameter — device assignment is deferred entirely to the scheduler. Validation still uses upstream exec_jobs (it runs on meta device, so no scheduling needed).

Behavioral change

max_workers now sets an upper bound on concurrency rather than a fixed thread count. Effective concurrency may be lower when GPU memory is tight. This also prevents OOMs caused by users picking too many workers.

TEST PLAN:

Unit tests (`test_scheduler.py`, no GPU required — runs on all CI):

estimate_job_memory: verifies sum-of-file-sizes × multiplier
_pick_device: most-free selection, returns None when nothing fits, respects reservation accounting, CPU always eligible
exec_jobs_dynamic CPU path: all jobs complete, empty input, order preservation

Integration tests (`test_multi_gpu.py`, `@requires_gpu(2)`):

test_multi_gpu_matches_single_gpu: saves Qwen3-0.6B with max_shard_size="1MB" (~198 shards), quantizes with FP8_dynamic on 1 GPU vs N GPUs, asserts bitwise tensor equality
test_multi_gpu_more_workers_than_shards: verifies graceful handling when max_workers > len(jobs)

Validated on:

2× NVIDIA RTX PRO 6000 Blackwell (homogeneous), Ubuntu 24.04, PyTorch 2.13.0+cu130
make quality passes (ruff check + ruff format + tools/lint_cuda.py)
All 11 tests pass

coderabbitai · 2026-07-28T15:47:33Z

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 8800efba-2d9f-46b4-a744-abcfa9d924b8

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

github-actions · 2026-07-28T15:47:34Z

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

mergify · 2026-07-28T15:48:01Z

Merge Protections

🔴 1 of 2 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require two reviews	👀 reviews
🟢	Require one maintainer review	—

🔴 Require two reviews

Waiting for

#approved-reviews-by >= 2

This rule is failing.

PRs labelled "two-reviews" must have at least two approving reviews before merging.

#approved-reviews-by >= 2
#changes-requested-reviews-by = 0

Show 1 satisfied protection

🟢 Require one maintainer review

All PRs must have at least one approving review from a maintainer before merging.

gemini-code-assist

Code Review

This pull request introduces a memory-aware dynamic job scheduler for multi-GPU model-free PTQ, replacing the previous static round-robin assignment with capacity-first scheduling based on real-time free VRAM. The reviewer provided valuable feedback to improve robustness and backward compatibility. Specifically, they recommended guarding device-module and cache-clearing operations against CPU devices and missing attributes to prevent runtime crashes, ensuring empty_cache() is called in the single-worker path to avoid false out-of-memory assumptions, and replacing torch.accelerator with torch.cuda in tests to maintain compatibility with PyTorch versions older than 2.4.

Replaces static round-robin device assignment with a capacity-first scheduler that checks real-time GPU memory before each job submission. Key changes: - New scheduler module with reservation tracking, caching allocator flush, and forced-fallback when nothing fits - model_free_ptq builds jobs without device assignment; scheduler picks GPUs dynamically at submit time Closes vllm-project#2975 Signed-off-by: Rohan Bandaru <rohanbanadaru14838@gmail.com>

Signed-off-by: Rohan Bandaru <rohanbanadaru14838@gmail.com>

kylesayrs

Nice job! Pinging @brian-dellabetta for review

Signed-off-by: Rohan Bandaru <rohanbanadaru14838@gmail.com>

brian-dellabetta

Hi @rohan9446 , thanks for preparing this. i have a handful of comments/questions, particularly whether we want this just in the model-free-ptq pathway and not the convert_checkpoint pathway in compressed-tensors. We have some work underway to consolidate the 2 entirely:

#2935

…uling, rename test Signed-off-by: Rohan Bandaru <rohanbanadaru14838@gmail.com>

kylesayrs

Nice work!

rohan9446 requested review from brian-dellabetta, dsikka and kylesayrs as code owners July 28, 2026 15:47

mergify Bot added the two-reviews When a PR requires two reviews label Jul 28, 2026

gemini-code-assist Bot reviewed Jul 28, 2026

View reviewed changes

rohan9446 force-pushed the feat/dynamic-gpu-scheduler branch from 15e596d to fd12f99 Compare July 28, 2026 16:01

rohan9446 force-pushed the feat/dynamic-gpu-scheduler branch from fd12f99 to d109158 Compare July 28, 2026 16:20

kylesayrs reviewed Jul 28, 2026

View reviewed changes

Rohan Bandaru added 2 commits July 28, 2026 20:19

address review: use torch.accelerator.memory APIs, simplify paths

285e65f

Signed-off-by: Rohan Bandaru <rohanbanadaru14838@gmail.com>

query memory once at startup, remove repeated empty_cache

9a8e11f

Signed-off-by: Rohan Bandaru <rohanbanadaru14838@gmail.com>

rohan9446 requested a review from kylesayrs July 29, 2026 03:56

Merge branch 'main' into feat/dynamic-gpu-scheduler

8bbc2a1

kylesayrs approved these changes Jul 30, 2026

View reviewed changes

Comment thread src/llmcompressor/entrypoints/model_free/scheduler.py Outdated

Rohan Bandaru and others added 2 commits July 30, 2026 12:52

remove unnecessary timeout from wait()

091720c

Signed-off-by: Rohan Bandaru <rohanbanadaru14838@gmail.com>

Merge branch 'main' into feat/dynamic-gpu-scheduler

6941274

kylesayrs approved these changes Jul 31, 2026

View reviewed changes

brian-dellabetta reviewed Aug 3, 2026

View reviewed changes

address brian review: dedup snapshot, consolidate with, fix CPU sched…

f0ba38f

…uling, rename test Signed-off-by: Rohan Bandaru <rohanbanadaru14838@gmail.com>

rohan9446 requested a review from brian-dellabetta August 3, 2026 18:16

Merge branch 'main' into feat/dynamic-gpu-scheduler

900dfd1

kylesayrs enabled auto-merge (squash) August 14, 2026 23:06

kylesayrs approved these changes Aug 14, 2026

View reviewed changes

Uh oh!

Conversation

rohan9446 commented Jul 28, 2026

SUMMARY:

Problem

Solution

Behavioral change

TEST PLAN:

Unit tests (test_scheduler.py, no GPU required — runs on all CI):

Integration tests (test_multi_gpu.py, @requires_gpu(2)):

Validated on:

Uh oh!

coderabbitai Bot commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

github-actions Bot commented Jul 28, 2026

Uh oh!

mergify Bot commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merge Protections

🔴 Require two reviews

🟢 Require one maintainer review

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kylesayrs left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

brian-dellabetta left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kylesayrs left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Unit tests (`test_scheduler.py`, no GPU required — runs on all CI):

Integration tests (`test_multi_gpu.py`, `@requires_gpu(2)`):

coderabbitai Bot commented Jul 28, 2026 •

edited

Loading

mergify Bot commented Jul 28, 2026 •

edited

Loading

brian-dellabetta left a comment •

edited

Loading