fix(loader): keep loader threads at one OpenMP thread on torch 2.11 by chelnnexy · Pull Request #2957 · ROCm/ATOM

chelnnexy · 2026-11-03T20:44:45Z

_one_omp_thread (#2871, made to resolve in #2936) is meant to give each weight-loader thread a one-thread OpenMP team. On torch 2.11 it doesn't hold. torch sets a thread's OpenMP count lazily, on that thread's first call into ATen, to the process default. The loader thread makes that first call after the initializer has run, so the 1 is overwritten and every loader thread runs the full default team (128 on a 256-core MI355X host).

tests/test_loader_threads.py::test_a_loader_thread_runs_one_omp_thread_and_no_other_thread_does fails on main for this reason:

>       assert pool.submit(torch.get_num_threads).result() == 1
E       assert 128 == 1

The fix

_one_omp_thread calls torch.get_num_threads() before omp_set_num_threads(1). That call is an ATen entry point, so it runs the thread's lazy initialization first, and nothing later resets the 1. The symbol is still resolved from torch._C at call time (as in #2936), and a failure raises, as on main.

Evidence, in a fresh thread on torch 2.11 (only one OpenMP runtime is mapped, libgomp.so.1, and the torch._C lookup from #2936 resolves into it):

Sequence in a new thread	`omp_get_max_threads()` / `torch.get_num_threads()`
`omp_set_num_threads(1)`, then read with OpenMP	1
...then the first `torch.get_num_threads()`	128 (torch reset it)
`torch.get_num_threads()` first, then `omp_set_num_threads(1)`	1

Tests

tests/test_loader_threads.py keeps main's two tests:

the setter is resolved from torch._C's library (opened == [torch._C.__file__]) and called with 1;
a real loader thread runs one OpenMP thread, also after a 512 × 512 matmul and a dtype round trip, and other threads keep the default. This is the test that fails on main.

Validation

CPU (CI environment: Python 3.12, CPU torch, .github/scripts/run_unit_tests.sh): 11115 passed, 0 failed. tests/test_loader_threads.py: 2 passed. black and ruff are clean on the changed files. CPU torch does not reproduce the bug (it keeps the 1), so the GPU host below is the real check.
MI355X (gfx950), torch 2.11.0+rocm10.0.0, on this head:
- tests/test_loader_threads.py: 20 of 20 runs passed (2 tests each); on main: 1 failed, 1 passed (128 == 1).
- Kimi-K3 + DSpark (TP8, 96 shards), weights loaded in, main and this PR alternately, two rounds:
  
  Run main 58ba501 this PR Change
  
  Round 1 119.5–119.9 s 94.3–94.7 s −21%
  
  Round 2 118.8–119.0 s 99.9–100.4 s −16%
  
  The range covers all 8 ranks. The server was healthy after 204 s and 202 s on main, and after 174 s and 188 s with this PR. Absolute load times vary between sessions on this host (an earlier session of the same comparison, on db16516, measured about 168 s on main and 88–118 s with the fix), so compare within a session.
- Each server answered 5 of 5 greedy smoke requests.

The other thread initializers in atom/ only set the CUDA device, so this pattern appears nowhere else.

`_one_omp_thread` sets the loader thread's OpenMP team to one thread, but torch initializes each thread's OpenMP count lazily, on the thread's first call into ATen, by setting it to the process default. On torch 2.11 that first call comes after the initializer, so it overwrote the 1 and every loader thread ran the full default team (128 on a 256-core MI355X host). `test_a_loader_thread_runs_one_omp_thread_and_no_other_thread_does` fails there with `128 == 1`. Call `torch.get_num_threads()` first, so torch's lazy initialization has already run when the count is set.

github-actions · 2026-11-03T20:45:04Z

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

✅ Run after the PR is approved and Pre Checkin passes
✅ Run immediately when an approval review is submitted
✅ Can be requested before approval with labels

Label	Tests
`ci:full`	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
`ci:atom`	Run native ATOM model accuracy tests
`ci:vllm`	Run ATOM vLLM OOT model accuracy tests
`ci:sglang`	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2957 --add-label <label>

…retry Improvements on the lazy-init fix (bb1c53a): Production code (loading_core.py): 1. Move ctypes.CDLL + symbol lookup to module scope so it runs once at import rather than once per loader thread (16 dlopen+dlsym -> 1). Guarded with try/except so a broken torch install does not make the entire module unimportable; the failure surfaces later with a clear warning when _one_omp_thread is actually called. 2. Replace the pre-trigger + set sequence with a set-verify-retry loop. The previous fix assumed torch.get_num_threads() is the specific call that triggers lazy init; if a future torch version triggers it on a different ATen entry point, the pre-trigger silently stops working. The loop doesn't care which call triggers it: the verify (torch.get_num_threads) is itself an ATen entry point, so if lazy init fires it will have fired by the time we check, and the retry lands on clean ground. If neither attempt succeeds, a warning surfaces the regression immediately. Tests (test_loader_threads.py): 3. main_thread_omp_guard fixture: _one_omp_thread is designed for worker threads. Unit tests that call it on the main thread previously left the main thread at 1 OMP thread, silently causing all later integration tests to pytest.skip. The fixture saves and unconditionally restores the count. 4. Split the single monkeypatch test into four: - succeeds_on_first_attempt: spy wraps real setter, loop exits in 1. - retries_when_first_set_is_overwritten: spy clobbers on first call (simulating lazy init), loop retries and succeeds. - warns_when_both_attempts_fail: no-op setter, verifies the warning. - survives_unresolved_symbol: _omp_set_num_threads is None, verifies graceful degradation. 5. test_omp_setting_survives_tensor_work: verifies the count stays 1 through a matmul and a dtype round-trip in a real thread pool.

The module-level lookup caught OSError and TypeError, but a torch whose OpenMP runtime has no omp_set_num_threads raises AttributeError, so loading_core failed to import and no model could load. Move the lookup into _resolve_omp_set_num_threads(), catch all three, and return None, which _one_omp_thread already handles with a warning. The main_thread_omp_guard fixture restored the count through loading_core._omp_set_num_threads, which the tests patch. It now keeps the real setter from before the test. Add a test that each of the three errors makes the resolver return None.

valarLip

Reviewed at db16516, by reading the diff and running the PR's tests (10 passed). This box has torch 2.9.1, where setting 1 and reading it back stays 1 (/app/logs_claude/pr2957_review/omp_lazy_probe.py), so I can't reproduce the 2.11 overwrite and I'm relying on the description's numbers for it.

The root cause is right. lazy_init_num_threads() uses a thread_local bool init, so a thread's first ATen call resets its OpenMP count to the process default, wiping the 1 the loader set. The fix is your first commit (bb1c53a): touch torch.get_num_threads() so the lazy init has run, then omp_set_num_threads(1). The next two commits build a retry loop, import-time resolution with swallowed errors, and tests for both on top of that. The third commit fixes an import failure the second one introduced. I'd like the PR brought back to bb1c53a's shape. Each item ends with the change I'm asking for.

Production code

1. The set-verify-retry loop adds nothing, and its docstring overstates it (loading_core.py:141-150). Lazy init happens once per thread. The loop works only because the verifying torch.get_num_threads() triggers that init, which is the same fact bb1c53a relies on, so the second iteration matters only on the first call and is dead afterwards. The docstring's "handles every ordering" doesn't hold. If a later torch stops initializing in get_num_threads, the first check reads 1 and returns, and a later ATen call resets it to 128 with neither the loop nor the warning noticing.
→ Change: two lines, torch.get_num_threads() (with one constraint comment: torch's per-thread lazy init must run before we set it) then omp_set_num_threads(1). Delete the loop and its warning.

2. Import-time resolution swallows the failure the old code raised (loading_core.py:103-118). _resolve_omp_set_num_threads catches AttributeError / OSError / TypeError and returns None, and _one_omp_thread then only warns. A torch without the symbol, or a torch._C that won't open, now silently falls back to the 128-thread team, which is the 168 s row in your own table, a 30–48% slower load, with one warning line. The stated reason ("16 dlopens become 1") doesn't hold up: re-opening an already-loaded library only bumps a refcount. And the AttributeError handling in the third commit is there because the second moved resolution to import time.
→ Change: resolve at call time from torch._C.__file__ as before, and let a failure raise. Delete _resolve_omp_set_num_threads, the module-level _omp_set_num_threads and the None branch.

3. One warning per loader thread (loading_core.py:136-140). _one_omp_thread is the pool initializer, so an unresolved symbol prints the same warning ATOM_LOADER_NUM_THREADS times on each of 8 ranks per load. The both-attempts-failed path does the same and calls torch.get_num_threads() a third time.
→ Change: goes away with item 2.

4. The same rationale three times, once wrongly (loading_core.py:92-102, the _resolve_* docstring, the _one_omp_thread docstring). The banner says a failure "surfaces later ... with a clear message", but it's a warning and the load continues. _one_omp_thread's docstring repeats the "each worker copies one tensor at a time" reason already given at its call site (~L261).
→ Change: keep one constraint comment in _one_omp_thread (why torch's thread init must run first) and delete the banner and the repeated paragraphs.

Tests

5. The assertion that guards the runtime choice was removed (tests/test_loader_threads.py:41). The old test asserted opened == [torch._C.__file__]. The new test_omp_resolved_from_torch_extension checks only not None and callable, though its name says "same library as torch". If someone switched to CDLL(None) or an absolute libgomp path, which can load a second OpenMP runtime (what #2936 fixed), it would still pass.
→ Change: restore the opened == [torch._C.__file__] assertion.

6. The unit tests cover only the parts items 1–2 remove (tests/test_loader_threads.py:72-130). retries_when_first_set_is_overwritten, warns_when_both_attempts_fail, survives_unresolved_symbol, the parametrized resolve test and the main_thread_omp_guard fixture all run on the main thread, which finished its lazy init long ago, so they can't reach "a new thread's first ATen call overwrote it". The retry test simulates the overwrite by having the spy write the default itself. The test that catches the real bug is the existing integration test, which fails on main with 128 == 1.
→ Change: delete these tests and the fixture (~80 lines) along with items 1–2.

7. test_omp_setting_survives_tensor_work duplicates the integration test (tests/test_loader_threads.py:137): same pool, same initializer, same default < 2 skip, same main-thread assertion. Checking that 1 survives a matmul and dtype round-trips is worth having.
→ Change: add one _work_then_report submit to test_a_loader_thread_runs_one_omp_thread_and_no_other_thread_does instead of a second function.

8. Section banners (tests/test_loader_threads.py:18, # ---- Fixture / Unit tests / Integration tests ----). They carry no constraint, and items 6–7 leave nothing to separate.
→ Change: delete them.

So the PR would be: bb1c53a's _one_omp_thread (touch get_num_threads, resolve from torch._C at call time, set 1) with one constraint comment, plus the original path assertion and the integration test with a tensor-work submit.

Per review: the fix is that torch.get_num_threads() runs the thread's lazy OpenMP initialization before omp_set_num_threads(1). The set-verify-retry loop only worked because its own check ran that initialization, so remove it. Resolve the symbol from torch._C at call time again and let a failure raise instead of warning and loading with the default thread count. In the tests, restore the check that the symbol comes from torch._C's library, drop the unit tests of the removed code and their fixture, and add a matmul and a dtype round trip to the integration test.

valar-cop

some changes needed..

green ci on db16516? all four workflows are cancelled

ci green.

valar-cop · 2026-11-04T02:41:22Z

+    return setter
+
+
+_omp_set_num_threads = _resolve_omp_set_num_threads()


returning None makes an unresolvable symbol a silent slow-path fallback, no clear message as the banner says. why not resolve at call time and let it raise?

Agreed. 7ddba25 removes _resolve_omp_set_num_threads and the None branch. The symbol is resolved from torch._C.__file__ at call time again, and a failure raises, as on main.

valar-cop · 2026-11-04T02:41:22Z

+            "loader threads will use the process-wide OpenMP default"
+        )
+        return
+    for _attempt in range(2):


the loop adds nothing over get_num_threads() then _omp_set_num_threads(1), both still depend on the same ATen call. maybe just drop the loop and keep one comment?

Agreed, the loop is gone in 7ddba25. _one_omp_thread now calls torch.get_num_threads(), then omp_set_num_threads(1), with one comment on why the order matters.

valar-cop · 2026-11-04T02:41:22Z

+# ---------------------------------------------------------------------------
+# Unit tests (mock / monkeypatch, run on the main thread)
+# ---------------------------------------------------------------------------
+def test_omp_resolved_from_torch_extension():


old test asserted the opened path is torch._C.file, this only checks not-None. CDLL(None) would still pass..

Right. 7ddba25 restores main's test, including assert opened == [torch._C.__file__], so CDLL(None) or another library path fails it.

valar-cop · 2026-11-04T02:41:22Z

-    def open_library(path):
-        opened.append(path)
-        return Runtime()
+def test_one_omp_thread_retries_when_first_set_is_overwritten(


these unit tests exercise the mock, not the defect, and one asserts the silent fallback as expected. remove them with the loop?

Agreed. 7ddba25 removes them, with the main_thread_omp_guard fixture. The integration test is the one that catches the defect: it fails on main on torch 2.11 with 128 == 1.

valar-cop · 2026-11-04T02:41:22Z

        assert torch.get_num_threads() == default
+
+
+def test_omp_setting_survives_tensor_work():


same pool and skip as the integration test, only the matmul is new. merge into one test?

Done in 7ddba25. The separate test is gone. The integration test now also submits _work_then_report (a 512 × 512 matmul and a dtype round trip) and checks that the count is still 1.

valar-cop · 2026-11-04T02:41:22Z

    )


+# ---------------------------------------------------------------------------


same rationale stated in the banner and the docstring, and the call-site comment repeats _one_omp_thread. keep one constraint comment..

Done in 7ddba25. The banner and the repeated paragraphs are gone. The docstring is main's, plus one constraint comment above torch.get_num_threads().

valar-cop · 2026-11-04T03:12:10Z

ci triage for db16516:

ATOM vLLM Test passed. https://github.com/ROCm/ATOM/actions/runs/41762379747

ATOM SGLang Test passed. https://github.com/ROCm/ATOM/actions/runs/41762379841

ATOM Test passed. https://github.com/ROCm/ATOM/actions/runs/41762380032

Pre Checkin passed. https://github.com/ROCm/ATOM/actions/runs/41741513939

zufayu

Reviewed the diff; approach is sound and the evidence table in the PR body (torch 2.11 lazy-init firing on the verify read, then settling) covers the mechanism. Set-verify-retry plus the integration test that does real ATen work after the initializer is the right regression net. Approving; a few non-blocking observations for follow-up, no changes requested:

test_omp_resolved_from_torch_extension hard-asserts the symbol is not None, and the spy tests (_spy/_clobber_once) would crash with TypeError rather than fail/skip when _omp_set_num_threads is None — yet production treats None as a supported, tolerated outcome. Minor inconsistency; only matters on torch builds where resolution legitimately returns None.
Docstring maps torch._C.__file__ is None to TypeError, but ctypes.CDLL(None) opens the main program — it would only raise AttributeError (or, worst case, mis-bind a foreign globally-loaded runtime; even then the verify fails and it degrades to warn+default, same as today). An explicit if torch._C.__file__ is None: return None guard would make the doc truthful.
The online-quant finalization pool (online_quant_streaming.py _worker_init) sets only the CUDA device/stream — consistent with the PR's claim that no other initializer needs this. Its workers run GPU-stream quant so OMP team size likely doesn't matter there, but worth a one-line note either here or there that the omission is deliberate.
Both failure warnings log once per worker per load (≈16 lines/rank × TP ranks). A log-once flag would keep the signal readable if the workaround ever regresses in the wild.

Perf numbers (−30%/−48% on TP8 weight loading) are a solid win.

chelnnexy · 2026-11-04T03:36:23Z

Thanks for the careful review. I agree with all eight points. 7ddba25 brings the PR back to bb1c53a's shape; against main it is now +3 lines in loading_core.py and +8 in the test.

Retry loop. Removed, with its warning. _one_omp_thread calls torch.get_num_threads(), then omp_set_num_threads(1), with one comment on why the order matters.
Import-time resolution. Removed: _resolve_omp_set_num_threads, the module-level _omp_set_num_threads and the None branch. The symbol is resolved from torch._C.__file__ at call time again, and a failure raises, as on main.
One warning per thread. Gone with item 2.
Repeated rationale. The banner and the repeated paragraphs are gone. The docstring is main's, plus the one constraint comment.
Path assertion. opened == [torch._C.__file__] is back (main's test_loader_thread_resolves_omp_from_torch_extension, unchanged).
Unit tests of the removed code. Deleted, with the main_thread_omp_guard fixture.
Duplicate integration test. Deleted. test_a_loader_thread_runs_one_omp_thread_and_no_other_thread_does now also submits _work_then_report (a 512 × 512 matmul and a dtype round trip) and checks the count is still 1.
Section banners. Gone; the file is main's plus that one submit.

On your torch 2.9.1 point: CPU torch 2.15 doesn't reproduce the overwrite either, so I re-ran the check on MI355X with torch 2.11.0+rocm10.0.0, at 7ddba25:

tests/test_loader_threads.py: 20 of 20 runs passed; main: 128 == 1, as before.
Kimi-K3 + DSpark, weights loaded in across the 8 ranks, two rounds:

Run main 58ba501 7ddba25 Change

Round 1 119.5–119.9 s 94.3–94.7 s −21%

Round 2 118.8–119.0 s 99.9–100.4 s −16%

The range covers all 8 ranks. The server was healthy after 204 s and 202 s on main, and after 174 s and 188 s with 7ddba25. Absolute load times vary between sessions on this host (an earlier session of the same comparison, on db16516, measured about 168 s on main and 88–118 s with the fix), so compare within a session.
CPU CI suite: 11115 passed, 0 failed.

The description is updated to match.

CI on db16516: it did pass. The push and the ready-for-review event each started all four workflows at the same time, and the concurrency group cancelled one copy of each. Those are the cancelled runs. The other copies passed:

Pre Checkin (run): Black, Ruff, catalog validation and the non-GPU unit tests.
ATOM Test (run): the offline inference smoke test. Accuracy was skipped by the heavy-CI gate.
ATOM vLLM Test and ATOM SGLang Test: success.

7ddba25 (the review changes, pushed now) starts a new set; I'll check that it goes green.

chelnnexy force-pushed the fix/loader-omp-lazy-init branch 3 times, most recently from 0acef71 to 82f4f46 Compare November 3, 2026 21:04

chelnnexy force-pushed the fix/loader-omp-lazy-init branch 2 times, most recently from eb7da54 to cccd694 Compare November 3, 2026 21:38

chelnnexy marked this pull request as ready for review November 3, 2026 23:44

zufayu requested a review from jiayyu November 4, 2026 01:29

zufayu assigned jiayyu Nov 4, 2026

valarLip reviewed Nov 4, 2026

View reviewed changes

valar-cop reviewed Nov 4, 2026

View reviewed changes

zufayu previously approved these changes Nov 4, 2026

View reviewed changes

chelnnexy dismissed zufayu’s stale review via 7ddba25 November 4, 2026 03:36

zufayu approved these changes Nov 4, 2026

View reviewed changes

valarLip merged commit 16ebac2 into ROCm:main Nov 4, 2026
50 checks passed

Run	main `58ba501`	this PR	Change
Round 1	119.5–119.9 s	94.3–94.7 s	−21%
Round 2	118.8–119.0 s	99.9–100.4 s	−16%

		return setter


		_omp_set_num_threads = _resolve_omp_set_num_threads()

		assert torch.get_num_threads() == default


		def test_omp_setting_survives_tensor_work():

		)


		# ---------------------------------------------------------------------------

Conversation

chelnnexy commented Nov 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

The fix

Tests

Validation

Uh oh!

github-actions Bot commented Nov 3, 2026

🏷️ CI Guide

Uh oh!

valarLip left a comment

Choose a reason for hiding this comment

Production code

Tests

Uh oh!

valar-cop left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

valar-cop commented Nov 4, 2026

Uh oh!

zufayu left a comment

Choose a reason for hiding this comment

Uh oh!

chelnnexy commented Nov 4, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

chelnnexy commented Nov 3, 2026 •

edited

Loading