TP > 1: every rank keeps torch's default intra-op thread count, so a multi-rank weight load busy-waits itself to a standstill on a small-core host (19 min -> 61 s)

[lukascechovic]

Summary

Nothing in python/freetoken/ sets torch.set_num_threads() on the load path, so every spawned TP
rank (server/launch.py:165) inherits torch's default of one intra-op thread per logical CPU. At
tp_size = N that puts N x cpu_count OpenMP threads on the host's physical cores, and libgomp's
default OMP_WAIT_POLICY=ACTIVE makes every idle thread busy-wait at each fork/join.

On a 6-core host at --tp-size 2 that is 24 threads on 6 cores, and each ~400 KB Tensor.copy_ in
the expert placement loop costs ~8.4 ms instead of ~0.04 ms. The cold load takes 19 minutes
instead of 61 seconds. The disk is idle throughout and both ranks sit at ~600 % CPU.

The trigger is tp_size x cpu_count > physical cores, so it is a property of the host shape, not of
any vendor or model. A large-core host at TP=2 will likely never see it.

The measurement, part 1: the mechanism, CPU only

A standalone harness, no GPU passed to the container, no model loaded, no server started. It imports
freetoken.models.weight.iter_expert_tensors_parallel and replays the six placement writes of
models/nvfp4_banks.py::load_nvfp4_expert_source_banks_parallel verbatim, in two spawn children
started the way server/launch.py starts ranks. Checkpoint is Qwen3.8-Flash-Next NVFP4: 193 expert
shards, 63.3 GiB, H=2560 I=640 E=512 L=48.

96 of the 193 expert shards, 31.66 GiB per rank, two ranks:

arm	rank wall	rate
as shipped (12 threads/rank, ACTIVE)	604.84 s	54 MiB/s
OMP_WAIT_POLICY=PASSIVE	15.94 s	2034 MiB/s (37.9x)
torch.set_num_threads(3)	12.82 s	2529 MiB/s (47.2x)

Thread-count sweep, 20 shards, two ranks:

threads per rank	rank wall
1	3.98 s
2	3.17 s
3	2.98 s
4	3.00 s
6	3.55 s
12 (the default)	131.85 s

Every count from 1 to 6 removes the collapse; only the default triggers it. The cliff is exactly at
2 ranks x 12 > 12 logical CPUs.

The tell that it is barrier overhead and not copy cost: with the default, all six placement
writes cost ~21.5 s each over 2560 calls — ~8.4 ms per call — and the write that moves one eighth
the bytes costs the same as the one that moves eight eighths. A copy whose cost is independent of its
size is not a copy cost.

What it is not. Two concurrent ranks move 13.19 GiB with 1.82 s each inside preadv (~7.2 GB/s
aggregate), so the disk is not the term. One process running the same placement shape finishes in
2.74 s, so the shape is not the term either. Only two ranks at once with the default thread count
collapses.

PASSIVE costs the single-rank case nothing: 2.76 s with it against 2.94 s without, one process, 20
shards.

The measurement, part 2: the real cold load, both GPUs

One env line on the container, OMP_WAIT_POLICY=PASSIVE. Same image, same command, no rebuild, no
code change. Control is the same box's own load ten hours earlier.

term	as shipped	PASSIVE
startup -> dense weights -> PLE	7 s	7 s
expert read	1147 s	32 s
pin settle (exposed once the read shrinks)	3 s	3 s
CUDA-graph capture	18 s	18 s
relay -> ready	0 s	1 s
load to ready	1175 s	61 s

The engine's own expert progress bar: 63.3G/63.3G [00:32<00:00, 2.48GB/s] against the control's
[19:07<00:00, 59.7MB/s].

Decode afterwards is unchanged — 35.2 / 35.2 / 35.7 / 35.7 tok/s, the same band as before the change
— so PASSIVE is not paid back at serve time even though it is process-wide and applies to every
ATen CPU parallel region for the server's life.

Suggested fix

server/launch.py::_run_scheduler, after import torch and before Scheduler(args):

if args.tp_info.size > 1:
    clamp = max(1, len(physical_core_cpus()) // args.tp_info.size)
    if clamp < torch.get_num_threads():
        logger.info_rank0(
            f"torch intra-op threads: {torch.get_num_threads()} -> {clamp} "
            f"(tp_size={args.tp_info.size})"
        )
        torch.set_num_threads(clamp)

physical_core_cpus() already exists at moe/cpu_executor.py:92, and cpu_executor.py:246-253
already does exactly this clamp-and-log — it just runs after the expert load, so it does not help
here. The two compose: the later clamp only ever lowers the count.

Measured, the clamp is slightly better than the env var (47.2x against 37.9x), which on our load is
worth ~6 s out of 61. The reason to prefer it is not the 6 s — it is that it is tp_size-aware and
protects every host, where an env var protects one deployment.

Happy to open this as a PR if you would like it in that form.

Scope, and what we did not measure

• Measured on our tree, which carries local patches: an intermediate-dim expert shard for TP
  (that is Offloaded MoE ignores tensor parallelism: TP=2 is a regression, and the fix is not expert parallelism #62's subject), feat(qwen4_exp): stream the PLE n-gram table from disk #311's disk PLE, and a chunked-multimodal-prefill patch. The mechanism
  does not depend on any of them — the harness imports your reader and replays your placement
  writes, and the trigger is only "N spawned ranks each defaulting to cpu_count intra-op threads".
• On upstream main each TP rank loads the full expert set rather than half, so if anything this
  should be worse there than here, not better. We have not run upstream main at TP=2 to confirm
  that and are not claiming it.
• Measured on the offloaded-MoE expert-bank load path. Any ATen CPU parallel region during a
  multi-rank load should be exposed the same way; we only measured this one.
• No quality claim — this box has no fidelity instrument, and nothing here touches numerics.
• Possibly adjacent, explicitly not claimed as the same bug: ft serve hangs at 96% weight load on Threadripper 1950X + RTX 5060 Ti #261 (ft serve hangs at 96 %
  weight load, load average > 10, box unresponsive) is single-rank, so this mechanism does not apply.
  But "CPU pegged, disk idle, machine unresponsive" is the same signature, and
  OMP_WAIT_POLICY=PASSIVE is a free thing to ask that reporter to try.

Platform

AMD Ryzen 5 7400 (6 physical cores / 12 logical) · AMD Radeon AI PRO R9700 x2 (gfx1201),
Linux, 122 GiB host RAM · PyTorch 2.11.0+rocm7.14.0, HIP 7.14.60850, RCCL 2.30.4 · FreeToken
main@4b94bdc3 + #132/#133/#134 + #311 + our sharding work.

The core count is the whole story here; the GPUs are incidental.