Repository navigation
Resident device AdamW: intermittent RNGD out-of-memory (DRAM_SHARED) after 3-13 steps #40
Description
Problem
Full FT with --resident-layers N (fp32 master + AdamW state on the card, optim_device.ResidentAdamW) stops with a device out-of-memory error after a few steps on RNGD, at sizes that should fit, and not every time:
furiosa_torch::cpp::allocator: Failed to alloc dram shared, required: 8388608, MemoryError(Allocate(8388608, 0, out of memory, ... capacity: 0
torch.OutOfMemoryError: failed to allocate 8388608 bytes on DRAM_SHARED
It fails in a host→device upload (FuriosaBackend.to_device) of an offloaded layer's new bf16 weight, from the layer worker thread.
Evidence
Llama-3.1-8B, S=512, accum 4, furiosa-torch 2026.2.0, driver/firmware 2026.4.0. Planned device memory per card is below the 47.5 GiB capacity in every case (bf16 weights + accumulated gradients of the card's layers, plus 2.62 GB per resident layer).
| run | resident layers per card | planned memory per card | steps before OOM |
|---|---|---|---|
| 1 card | 8 | ~50 GB | 0 (first step; really too big) |
| 1 card | 4 | ~40 GB | 12 |
| 1 card | 3 | ~37 GB | none (25/25) |
| 2 cards | 8 | ~36 GB | 3, 7, 13, and once none (25/25) |
| 2 cards | 6 | ~31 GB | 7 |
| 2 cards, offload only | 0 | ~15 GB | none (25/25, twice) |
- The 2-card, 8-layer run reached step 3 before Full-FT offload is host-bound: measure resident device AdamW and update the placement model #34's fixes and step 7–13 after the PR for Full-FT offload is host-bound: measure resident device AdamW and update the placement model #34 frees the old bf16 weights before the device update. One run of the same code finished all 25 steps.
- In isolation the pieces do not leak: 400 calls of the [14336, 4096]
ResidentAdamWgraph with its state replaced each call, 400 [14336, 4096] bf16 uploads replaced each time, and 20 000 [1,1] scalar uploads all finish on one card. furiosa-smi statusreports 0.00/47.50 GiB for all cards throughout the run, so device memory use can't be watched from outside.- The host also ran another user's CPU-heavy job during the later runs (load ~100 on 96 cores), so the timing is not reproducible. Whether that affects the allocator is unknown.
Records: experiments/2026-10-21-resident-adamw.md (in the PR for #34).
Scope
- Find out what grows: fragmentation from freeing and re-allocating the fp32 state (w, m, v) every step while the worker threads upload bf16 weights, versus a leak in the runtime.
- If the runtime (furiosa-torch 2026.2) needs a fix, file it with Furiosa and keep a workaround here, for example fewer resident layers, or reusing output buffers if the runtime allows it.
Done when
train --full --resident-layers 8 --backend furiosa:auto,furiosa:auto runs 200 steps on two RNGD cards without OOM, or the limit is documented with its cause.
Owner
gspark-etri
Activity
Resident device AdamW: compile in 150 s, RNGD measurements, host-thro…
Record the DRAM_SHARED race (#40): probes, repro, six 25-step residen…
gspark-etri commentedon Oct 21, 2026
Root cause: a race on furiosa-torch's process-global memory-type flag, not running out of memory
What fails. furiosa-torch 2026.2 keeps the device memory type (DRAM or DRAM_SHARED) in one flag for the whole process (
furiosa.native_torch.memory._set_dram_shared_mode/_set_memory_type). Every compiled-graph call sets that flag to DRAM_SHARED for a moment. If another thread allocates a device tensor at that moment, the tensor goes into DRAM_SHARED. On a single-chip card that pool is 0 bytes (thecapacity: 0in the error), so the allocation fails at once. The full-FT layer workers upload bf16 weights while the main thread runs graphs, so they hit this.Why resident runs fail and offload-only runs don't. In resident runs the main thread runs the resident AdamW graphs (7 per resident layer) exactly while the workers start their uploads. In offload-only runs the main thread mostly waits on the upload futures, so the two overlap less. They still can, just less often.
Probes (1 RNGD card, furiosa-torch 2026.2.0, driver/firmware 2026.4.0):
... on DRAM, and graph outputs stop at 47.38 GiB. 8 GiB of uploads plus graph outputs stop at 39.38 GiB. Nothing normally goes into DRAM_SHARED._set_dram_shared_mode()called in a thread is seen by the main thread (_is_dram_shared_mode() == True). An 8 MB upload from the main thread then fails withfailed to allocate 8388608 bytes on DRAM_SHARED, the same error as here.Repro:
experiments/data/dram-shared/race_repro.pyin #42 (flag,watch,race,fixed).Workaround in npu-train (#42). Graph calls and device allocations now take one process-wide lock (
backends.DEVICE_ALLOC_LOCK). The host→device copy runs outside the lock, so upload speed is unchanged: 6.01 vs 6.05 GB/s for 117 MB. With the lock, the same race ran 120 s (58,799 graph calls, 60,725 uploads) with no errors.Result. After the fix, six back-to-back 25-step resident runs all finished with no OOM: 2 cards × 8 resident layers three times, and 1 card × 4 resident layers three times. Their losses match the earlier resident runs bit for bit. Host load was 350–650 during these runs, so their step times are not performance numbers.
For Furiosa: the memory type should be thread-local, or graph calls should not change it on single-chip devices. Please pass this on together with the repro.
Low-load resident timing for #34 and the step model; docs for the #40 …
Serialize RNGD device allocations with graph calls (DRAM_SHARED race, #…
Record the DRAM_SHARED race (#40): probes, repro, six 25-step residen…
Low-load resident timing for #34 and the step model; docs for the #40 …
Fix resident-AdamW OOM on DRAM_SHARED: serialize device allocations w…
gspark-etri commentedon Oct 21, 2026
PR #42 merged: root cause fixed (process-wide DRAM_SHARED flag raced by layer-worker uploads; one lock for graph calls and device allocations). 8/8 resident runs reached 25/25 steps without OOM. Staying open for the done-when 200-step 2-card resident run, to be done only when two cards are free and host load is low, as one run-and-exit Job.
Resident AdamW: 200 steps on 2 RNGD x 8 layers without OOM (#40) + 1-…