[Examples] Glm5.2 MXFP4xMXFP8 by kylesayrs · Pull Request #3048 · vllm-project/llm-compressor

kylesayrs · 2026-08-18T02:51:26Z

Prerequisites

feat: dynamic memory-aware GPU scheduling for model-free PTQ #2976

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 #2975 Signed-off-by: Rohan Bandaru <rohanbanadaru14838@gmail.com>

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

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

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

coderabbitai · 2026-08-18T02:51:34Z

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: cacfcad1-8f3b-410d-ae9b-0ea9b2a519d6

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

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.}

mergify · 2026-08-18T02:52:04Z

Merge Protections

🟢 Merge protection satisfied — ready to merge.

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 old round-robin assignment with capacity-first scheduling based on real-time VRAM estimates. It also adds corresponding unit and multi-GPU tests, alongside an experimental script for GLM-5.2 quantization. Feedback focuses on improving the scheduler's robustness by force-scheduling jobs when conservative memory estimates exceed available VRAM (rather than raising a RuntimeError), updating the associated unit tests, and correcting a non-existent model ID (Qwen/Qwen3-0.6B) used in the multi-GPU tests to an existing one like Qwen/Qwen2.5-0.5B.

gemini-code-assist · 2026-08-18T02:53:17Z

+            if not inflight:
+                if not pending:
+                    break
+                raise RuntimeError(
+                    "No device has enough estimated free memory "
+                    "for any remaining shard. Consider reducing "
+                    "max_workers or adjusting _MEMORY_MULTIPLIER."
+                )


Instead of raising a RuntimeError and crashing when a job's conservative memory estimate exceeds the available VRAM on all devices, it is much more robust to force-schedule the job on the device with the most available memory. This aligns with the single-worker path behavior (which warns and proceeds) and prevents hard crashes on conservative overestimates, allowing the quantization to proceed and potentially succeed.

Suggested change

if not inflight:

if not pending:

break

raise RuntimeError(

"No device has enough estimated free memory "

"for any remaining shard. Consider reducing "

"max_workers or adjusting _MEMORY_MULTIPLIER."

)

if not inflight:

if not pending:

break

idx = pending[0]

dev = max(devices, key=lambda d: _free_bytes(d, initial_free, reserved))

logger.warning(

f"Shard {idx} (~{memory_estimates[idx] / 1e9:.2f} GB) "

f"exceeds estimated available capacity on all devices. "

f"Force-scheduling on {dev} anyway."

)

fn, iwm, sp, sch, ign, conv = jobs[idx]

fut = pool.submit(fn, iwm, sp, sch, ign, dev, conv)

inflight[fut] = idx

fut_device[fut] = dev

reserved[dev] += memory_estimates[idx]

pending.remove(idx)

gemini-code-assist · 2026-08-18T02:53:17Z

+def test_multi_gpu_matches_single_gpu(tmp_path):
+    """Dynamic scheduling across N GPUs should give identical output to
+    a single-GPU run."""
+    model_id = "Qwen/Qwen3-0.6B"


The model ID Qwen/Qwen3-0.6B does not exist on Hugging Face (Qwen3 has not been released yet). This will cause the test to fail with a RepositoryNotFoundError. Please use an existing small model such as Qwen/Qwen2.5-0.5B or Qwen/Qwen2-0.5B.

Suggested change

model_id = "Qwen/Qwen3-0.6B"

model_id = "Qwen/Qwen2.5-0.5B"

gemini-code-assist · 2026-08-18T02:53:17Z

+@requires_gpu(2)
+def test_multi_gpu_more_workers_than_shards(tmp_path):
+    """Scheduler should be fine when there are more threads than shards."""
+    model_id = "Qwen/Qwen3-0.6B"


The model ID Qwen/Qwen3-0.6B does not exist on Hugging Face. Please use an existing small model such as Qwen/Qwen2.5-0.5B or Qwen/Qwen2-0.5B.

Suggested change

model_id = "Qwen/Qwen3-0.6B"

model_id = "Qwen/Qwen2.5-0.5B"

gemini-code-assist · 2026-08-18T02:53:17Z

+@patch(
+    "llmcompressor.entrypoints.model_free.scheduler"
+    ".torch.accelerator.memory.get_memory_info"
+)
+def test_raises_when_no_device_fits(mock_mem_info):
+    """Multi-worker path should raise RuntimeError when no device can fit
+    any remaining shard and nothing is in flight."""
+    mock_mem_info.return_value = (1000, 96_000_000_000)
+
+    def fn(iwm, sp, sch, ign, dev, conv):
+        return (0, {})
+
+    # job needs 10 GB but device only has 1000 bytes
+    jobs = [(fn, {}, "s0.st", "FP8", [], None)]
+    mem = [10_000_000_000]
+
+    with pytest.raises(RuntimeError, match="No device has enough"):
+        exec_jobs_dynamic(
+            jobs,
+            [torch.device("cuda:0")],
+            2,
+            mem,
+            desc="Test",
+        )


Update the unit test to verify the new robust force-scheduling behavior when no device has enough estimated free memory, rather than expecting a RuntimeError.

@patch( "llmcompressor.entrypoints.model_free.scheduler" ".torch.accelerator.memory.get_memory_info" ) def test_force_schedules_when_no_device_fits(mock_mem_info): """Multi-worker path should force-schedule when no device can fit any remaining shard and nothing is in flight, rather than raising RuntimeError.""" mock_mem_info.return_value = (1000, 96_000_000_000) def fn(iwm, sp, sch, ign, dev, conv): return (100, {"t": sp}) # job needs 10 GB but device only has 1000 bytes jobs = [(fn, {}, "s0.st", "FP8", [], None)] mem = [10_000_000_000] out = exec_jobs_dynamic( jobs, [torch.device("cuda:0")], 2, mem, desc="Test", ) assert out == [(100, {"t": "s0.st"})]

github-actions · 2026-08-18T02:53:36Z

👋 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-08-18T08:56:37Z

The quality checks have failed. Please run make style and make quality under
the root directory to adddress the lint failures. You will need to install the
dev optional install to get the required linting packages:
https://github.com/vllm-project/llm-compressor/blob/main/CONTRIBUTING.md

brian-dellabetta · 2026-08-19T15:00:57Z

+        r"re:.*mlp\.gate.*",
+        r"lm_head",
+    ],
+    max_workers=50,


isn't this rather high?

It's a large model, honestly afaict we can go larger without suffering too much overhead.

brian-dellabetta · 2026-08-19T15:01:30Z

+SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-MXFP4xMXFP8"
+
+# use `model_free_ptq` to apply quantization
+# NOTE: MXFP4xMXFP8 is experimentally supported in vllm


how experimental? 😅

vllm-project/vllm#51332

Rohan Bandaru and others added 11 commits July 28, 2026 09:20

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>

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

8bbc2a1

remove unnecessary timeout from wait()

091720c

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

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

6941274

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

f0ba38f

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

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

900dfd1

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

bd37b25

add example

dcffdd3

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

move

d1ecf4d

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

kylesayrs requested review from brian-dellabetta and dsikka as code owners August 18, 2026 02:51

gemini-code-assist Bot reviewed Aug 18, 2026

View reviewed changes

mergify Bot added quality-failed and removed quality-failed labels Aug 18, 2026

Merge branch 'main' into kylesayrs/glm52-mfptq

d1ced44

kylesayrs added the ready When a PR is ready for full CI testing before merge label Aug 19, 2026

brian-dellabetta approved these changes Aug 19, 2026

View reviewed changes

Merge branch 'main' into kylesayrs/glm52-mfptq

8f2c123

kylesayrs enabled auto-merge (squash) August 19, 2026 16:40

kylesayrs disabled auto-merge August 19, 2026 16:40

-            if not inflight:
-                if not pending:
-                    break
-                raise RuntimeError(
-                    "No device has enough estimated free memory "
-                    "for any remaining shard. Consider reducing "
-                    "max_workers or adjusting _MEMORY_MULTIPLIER."
-                )
+            if not inflight:
+                if not pending:
+                    break
+                idx = pending[0]
+                dev = max(devices, key=lambda d: _free_bytes(d, initial_free, reserved))
+                logger.warning(
+                    f"Shard {idx} (~{memory_estimates[idx] / 1e9:.2f} GB) "
+                    f"exceeds estimated available capacity on all devices. "
+                    f"Force-scheduling on {dev} anyway."
+                )
+                fn, iwm, sp, sch, ign, conv = jobs[idx]
+                fut = pool.submit(fn, iwm, sp, sch, ign, dev, conv)
+                inflight[fut] = idx
+                fut_device[fut] = dev
+                reserved[dev] += memory_estimates[idx]
+                pending.remove(idx)

Uh oh!

Conversation

kylesayrs commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Prerequisites

Uh oh!

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

Uh oh!

Review skipped

Uh oh!

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

Uh oh!

Merge Protections

🟢 Require one maintainer review

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Aug 18, 2026

Uh oh!

mergify Bot commented Aug 18, 2026

Uh oh!

brian-dellabetta Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

kylesayrs Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

brian-dellabetta Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

kylesayrs Aug 19, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kylesayrs commented Aug 18, 2026 •

edited

Loading

coderabbitai Bot commented Aug 18, 2026 •

edited

Loading

mergify Bot commented Aug 18, 2026 •

edited

Loading