Multi-host export compiles for 22 minutes, then fails, when one rank has no memory for a forward compiler #242

What happens

On three DGX Spark GB10 hosts, one rank per host, a flare export of DeepSeek V4.1 Flash with tensor parallelism 3 on d2b5c5f3 failed after "Compile model forwards" had already run for 1310.8 s:

[    OK    ] Compile model forwards.  38/38  1310.8s
RuntimeError: forward compilation failed: [(1, 'RuntimeError: MemoryError: preparation requires at least 12288 MiB available host memory; found 12038 MiB'), (2, 'RuntimeError: MemoryError: preparation requires at least 12288 MiB available host memory; found 11939 MiB')]

Each rank's forward compiler pool (CompilePool) sizes its workers with worker_limit (ignition/preparation/host_memory.py:105), which requires available memory of at least the larger of a fixed 8 GiB floor (_HEADROOM_BYTES) and the recipe's preparation headroom (1 GiB in this plan), plus 4 GiB per worker: 2 GiB of host memory (_WORKER_BYTES) and 2 GiB for the worker's CUDA context (_WORKER_DEVICE_BYTES, compile_pool.py). At that point in the export (weights, KV cache and graph operands resident, after kernel tuning), ranks 1 and 2 were 250 MiB and 349 MiB short of the 12288 MiB requirement for one worker, so neither could create a pool. Rank 0 had enough memory and compiled its own share of the forwards. The progress line counts only rank 0's share, and the phase is marked OK because the outcome check runs after the phase closes; the export then failed on the errors ranks 1 and 2 reported in the outcome all-gather at the end of the phase.

A later export with the same model and options on the same hosts passed this phase. Its plan sized 1099 more KV pages (128643 vs 127544), yet rank 0 reported 10329 MiB free before kernel tuning against 9610 MiB in the failing export, with identical allocator figures in both (100950 MiB allocated, 101056 MiB reserved). Kernel tuning in the passing export measured and compiled no kernels (0 measured, 0 compilations) and took 358 s; in the failing export it reported 809 measured and took 4492 s. On rank 0, CUDA-reported free memory before kernel tuning therefore differed by about 700 MiB, in the opposite direction from the KV plan and outside PyTorch's caching allocator. The pool check uses the host's available memory instead, which these logs do not record before tuning, and they record only rank 0; which allocation accounts for the difference was not measured.

Why it fails instead of degrading

_prepare_ranks (ignition/preparation/aot.py) already handles a rank without a compiler: the MemoryError is caught at line 803 and the comment says that rank "lends its forwards to the others on this host". Lending goes through the host's shared cache directory, so it is enabled only when the group does not span hosts (line 827). Here each host has one rank, so a rank without a compiler has no rank to lend to, and it fails at lines 855-857 as soon as it has a forward missing from its cache.

Two consequences:

1. A multi-host export fails when any rank cannot create a compiler and has at least one forward missing from its cache, even though another rank has a working pool and the cross-host path already sends compiled artifacts to the hosts that lack them (line 911).
2. The failure is known when the pools are created, and the all-gather at line 826 already gives every rank each rank's count of forwards missing from its cache and its compiler worker count, but ranks 1 and 2 report their failure only in the outcome all-gather at line 901, so the export stops after rank 0 has compiled its whole share; here that was about 22 minutes.

Possible directions

• Let a rank on another host compile the forwards of a rank without a compiler, sending the owner's forward metadata over the group instead of the shared directory, and then reuse the existing cross-host artifact transfer. The borrowing rank would then report the artifact in its own outcome, since _share_artifacts sends each file from the rank that lists it. In this case rank 0 would compile the other ranks' shares with its own workers, so the phase would take longer than when every rank compiles.
• If that is not wanted, fail right after the all-gather at line 826 on every rank when any rank has missing forwards, no workers, and cannot lend them, so the export stops before compiling.

PR #122 makes this 8 GiB floor configurable; setting it to 6 GiB would lower the requirement here to 10240 MiB (6 GiB floor plus one 4 GiB worker; the 1 GiB preparation headroom is smaller), which both hosts met (untested on this failure). That does not change how a multi-host group handles a rank without a compiler.

Environment

• Three DGX Spark GB10 hosts (unified memory), flare export with tensor parallelism 3, one rank per host.
• Checkpoint deepseek-ai/DeepSeek-V4.1-Flash-1017, recipe defaults, compile_workers unset.
• Flare d2b5c5f3.