Skip to content

Resident device AdamW: compile in 150 s, RNGD measurements, host-throughput step model - #41

Merged
gspark-etri merged 1 commit into
mainfrom
perf/34-resident-adamw-compile
Oct 21, 2026
Merged

gspark-etri merged 1 commit into
mainfrom
perf/34-resident-adamw-compile

Conversation

@gspark-etri

Copy link
Copy Markdown
Contributor

Background

#34: after #26, full FT is host-bound. The resident device AdamW (--resident-layers, from #30) had never been measured on RNGD because npu-lora compile --resident-layers 1 did not reach the AdamW graphs in 2.5 h. The §5.3 optimizer-placement model charged link bandwidth and predicted 5.6 s for a step that measured 9.26 s.

Why the compile was slow: cmd_compile runs the whole training setup first. It loads the 8B weights and the 32 GB fp32 master, traces the layer graphs (with a CPU check run), and compile_parts runs every part's sample on the CPU, even when all parts are cached. The 10-21 log shows the 11 forward parts loaded from cache in 0 s, then no output for more than 2 h inside the backward parts' CPU runs (shared node, 8 CPUs requested). The AdamW graphs depend only on the weight shapes.

Changes

  • npu-lora compile --adamw-only: reads the 2-D layer weight shapes from the safetensors header (Weights.shape, optim_device.resident_shapes) and compiles only the ResidentAdamW graphs. It loads no weights and builds no training graph. The trainer uses the same shape helper.
  • --backend furiosa:auto,furiosa:auto uses the pod's assigned cards in order. Before, on a 3-card node furiosa:0,furiosa:1 could open a card the pod was not assigned, and visible_chip read a multi-card file such as npu1,npu3 as one number (13).
  • train --update-workers N (default 4) sets the number of host threads for the per-layer offload jobs.
  • Resident update: the old bf16 device weights are dropped before the update (as _layer_job already does), and each layer's accumulated gradients are freed right after it.
  • docs/design.md: new section "Optimizer placement: step-time model", stated in host copy + host AdamW throughput terms. Records: new experiments/2026-10-21-resident-adamw.md (KO, like the other records), a follow-up note in 2026-10-21-two-card-split.md, raw step logs in experiments/data/resident-adamw/. README EN/KO: --adamw-only. STATUS.md/.research/events.jsonl via report.py.

Verification

  • CPU: python tests/test_cpu.py with torch 2.9.1+cpu and transformers 5.1.0 gives 23/23 ok. New tests: test_resident_shapes_from_header and test_visible_chip_picks_the_nth_assigned_card. test_two_device_split_and_resident_adamw_equal_one_device still passes.
  • python .research/check_sensitive.py reports clean.
  • RNGD: npu-node-1, furiosa-torch 2026.2.0, torch 2.10.0+cpu, driver/firmware 2026.4.0, furiosaai/furiosa-llm:latest. Before the runs, no process on the node held /dev/rngd*. I used at most 2 cards and did not use npu-node-2. Recipe from the two-card record: Llama-3.1-8B-Instruct, S=512, accum 4, lr 1e-5, nsmc_mix, seed 0, 25 steps, median of steps 6–25.
    • Compile: compile --adamw-only took 150 s on a CPU node: 4 graphs, at 4, 15, 58 and 73 s. Training runs load these graphs from the cache.

    • Runs before host contention (13:15–14:11 KST):

      run step tok/s update wait
      1 card, offload 5.22 s 392 1.50 s
      1 card, offload, --update-workers 8 4.82 s (−8 %) 425 1.23 s
      2 cards, offload 5.09 s 402 1.50 s
      2 cards, 8 resident layers per card (steps 3–7; OOM at 7) 4.57 s (+14 % tok/s) 448 0.15 s
      1 card, 4 resident layers (steps 6–12; OOM at 12) 5.12 s 400 1.04 s
      1 card, 8 resident layers OOM at step 1 (~50 GB, too big for one card)
    • Loss: the offload runs agree bit for bit (1.7157, 1.5626, 1.0879, 1.1060, 1.3361). The resident runs give the same step 1 and differ by at most 0.5 % at steps 2–5 (2 cards × 8: 1.5649, 1.0858, 1.1035, 1.3369).

    • Under host contention: another user's CPU job ran on the same node from about 14:15 (load ~100 on 96 cores). Back to back with this PR's code: 2-card offload 10.96 s, 2-card × 8 resident 8.83 s (−19 %; steps 6–7, OOM at 7).

    • Model check: within 15 % on 6 configurations (−10 % to +10 %). For example, 2 cards × 8 resident: predicted 4.33 s, measured 4.57 s. For the 24-core 2-card offload step from the earlier record: predicted 8.33 s, measured 9.26 s; the link model had predicted 5.6 s (−40 %). Full table in design.md and the record.

  • How to reproduce: the commands are in the record. Jobs were built with the experiments/data/w7/jobgen.py pattern. All Jobs, pods and ConfigMaps are deleted.

Compatibility and limits

  • The resident runs hit an intermittent device OOM (DRAM_SHARED) after 3–13 steps, even when the planned memory fits. One 2-card × 8 run finished all 25 steps, and a 1-card × 3 run finished too. Repeating each piece on its own (the AdamW graph, big uploads, scalar uploads) does not grow memory. The cause is unknown, so I filed Resident device AdamW: intermittent RNGD out-of-memory (DRAM_SHARED) after 3-13 steps #40. That is why the resident step times above cover only 2–7 steps rather than 6–25. Until Resident device AdamW: intermittent RNGD out-of-memory (DRAM_SHARED) after 3-13 steps #40 is fixed, --resident-layers is not usable for real training runs.
  • I tried running the resident updates before the offload uploads, to test a fragmentation hypothesis. The results were mixed, so I reverted it. The runs that used that variant are marked in the record.
  • Everything new is opt-in, and the defaults are unchanged (--update-workers 4, furiosa:N as before).

Related Issue

Related to #34 (compile fix, measurements and model are done; full 6–25-step resident medians wait on #40). Related to #40.

🤖 Generated with Claude Code

…d resident-layer measurements (#34)

- npu-lora compile --adamw-only: compile the ResidentAdamW graphs from the
  safetensors header shapes, without weights or training graphs (150 s
  for Llama-8B, was >2.5 h through the full compile path)
- furiosa:auto,furiosa:auto picks the pod's assigned cards in order
  (multi-card file 'npu1,npu3' was read as one number)
- --update-workers N for the host offload threads (8: -8 % step)
- free old bf16 weights and gradients before/after the device update
- docs/design.md: optimizer placement step-time model in host terms;
  experiments/2026-10-21-resident-adamw.md with the RNGD runs

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@gspark-etri
gspark-etri merged commit a07aefc into main Oct 21, 2026
1 check passed
@gspark-etri
gspark-etri deleted the perf/34-resident-adamw-compile branch October 21, 2026 06:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant