Repository navigation
Resident AdamW: 200 steps on 2 RNGD x 8 layers without OOM (#40) + 1-card x 4 low-load timing (#34) - #45
Merged
Conversation
This was referenced Oct 21, 2026
Closed
gspark-etri
added a commit
that referenced
this pull request
Oct 21, 2026
## Background npu-lora could not train on ATOM (#15). The first version of this PR compiled 50 of 51 graph parts with rebel-compiler 0.11.0, but it left two problems: - one backward part (the SiLU backward) crashed the compiler, so it ran on the host; - `npu-lora check --backend rbln:0` failed: the worst LoRA gradient was 16 % off the CPU, against a 5 % threshold. This update finds the cause of both and makes `check` pass. The branch is rebased on `main` (#42–#45). ## Changes - `npu_lora/backends.py` `RBLNBackend` (`--backend rbln:N`): - Graph parts are compiled with `rebel.compile_from_torch` and run by an RBLN runtime, with host tensors in and out. - Compiled parts are cached as `.rbln`. Each part compiles in a child process, and a part that fails to compile falls back to the host. - Arguments are mapped by name to the inputs the compiled graph kept. - **New**: before compiling, `graph.fold_neg_add` rewrites `add(neg(x), c)` → `rsub(x, c)`. rebel-compiler 0.11 rejects that pair (DEVICE_GRAPH_CONVERSION), and it segfaults when a matmul follows; this was the SiLU backward. The cache path hashes the rewritten graph. - **New**: `host_parts` / `--host-parts` puts whole part groups (`fwd`, `bwd`, `rms`, `ce-stats`, `ce-grad`, `rms-bwd`) on the host on purpose. - Default for rbln: `fwd`. The layer forward runs on the host; the layer backward (which recomputes the forward) and the head run on ATOM. - `none` puts every part on ATOM. - Steps report `parts_device_s` / `parts_host_s`. Log lines go to stderr. - `npu_lora/cli.py`: - `check` also reports both runs against an fp32 autograd reference (`--reference fp32`, dense LoRA). The pass criterion is unchanged: NPU vs CPU. - `_backend` passes `--cache-dir` and `--host-parts` to rbln. - `npu_lora/trainer.py`: `LoRATrainer.reference_grads` (fp32 autograd with the graphs' bf16 weights and adapters). - Tests (CPU): `fold_neg_add` keeps the SiLU-backward values; host placement and group parsing; the fp32 reference matches HF's loss and the CPU gradients. - `experiments/2026-10-22-rbln-atom-check.md` + `experiments/data/rbln-atom/`: the record, raw results, and two standalone repros for the vendor. - README / README.ko: the backend table and Limitations. ## Verification - Hardware and SDK: - Host: npu-node-1, one RBLN-CA22 used directly; each run used a card with no process in `rbln-stat`. - rebel-compiler 0.11.0, torch 2.11.0+cpu, transformers 5.8.1. - Newest SDK found: 0.12.0. optimum-rbln, vllm-rbln and torch-rbln 0.12.0 are on PyPI (2026-10-16) and need rebel-compiler 0.12.0. That is only on the vendor index, which answers 401 (no credentials on the node), so it was not tested. - Inputs: - `check`: `unsloth/Llama-3.2-1B-Instruct`, first 2 layers, S=128, LoRA r16 q·v, B ~ N(0, 0.01), random ids (check's seed). The diagnostics also varied the ids and B (seeds 2–4). - Training: all 16 layers. - Criteria: `check` passes when the loss is within 1 % and every LoRA gradient within 5 % of the CPU (unchanged), and a LoRA run reduces the loss. - Method: before choosing a placement, every ATOM part call was also run on the CPU with the same inputs, and parts were moved to the host by group. - Results: - CPU suite on Windows / torch 2.9.1: the 3 new or extended tests pass. 8 tests also fail on `main` in that local environment (transformers 4.57, ai-rm entry points not installed); CI runs the full suite. - Compile: before, 1 of 39 backward parts was on the host (segfault). After, all 59 parts compile. The SiLU-backward part is within 0.7 % of the CPU on ATOM. - `npu-lora check --backend rbln:0 --layers 2 --seq-len 128 --max-mm 2`: | placement | loss rel. | worst grad vs CPU | min cos | pass | |---|---:|---:|---:|---| | `--host-parts none` (59 parts on ATOM) | 1.76e-3 | 13.36 % | 0.99114 | no | | default (`fwd` on the host: 44 parts on ATOM, 15 on the host) | 1.84e-3 | **1.39 %** | 0.99991 | **yes** | The default was re-run from the pushed head (`a472883`): same numbers, exit code 0. - Other inputs (seeds 2–4), worst gradient vs CPU: - default placement: 4.9 / 1.7 / 2.7 %; - all on ATOM: 10.1 / 2.2 / 3.1 %. - Cause: ATOM rounds away from zero far more often than round-to-nearest-even (bias +0.5…+0.9), and fp32 outputs have bf16-level error (repro table in the record). - Every part alone is within 0.75 % of the CPU (`ce-grad` up to 4.3 %), but the forward activations that leave a layer drift and the head's softmax amplifies that. - Most of the effect comes from the q/k/v projection parts. - A head on the host does not help (13.4 %). Splitting the head (`ce-grad` on the host, `lse` from ATOM) makes it worse. - The fp32 reference shows the CPU bf16 run is not exact either: on check's inputs it is 9.1 % off fp32. On seeds 2–4 all-ATOM is 1.2–3× further from fp32 than the CPU. - LoRA, r16 q·v, S=128, default placement vs `--backend cpu`, same code and seed: - code-name set (64 examples), 30 steps, lr 5e-4: ATOM 3.85 → 0.54, CPU 3.81 → 0.51. - NSMC train, 60 steps, lr 2e-4, mean of steps 1–10 → 51–60: - ATOM default 0.470 → 0.094; - all on ATOM 0.479 → 0.128; - CPU 0.448 → 0.093. - The curves track the CPU (median 10-step window difference 4.9 %). - The reported ATOM loss sits about 0.035 nats above the CPU on confident tokens, because `lse` comes from ATOM's logits and the target logit from the host. - Median step time, 1B: ATOM 1.39 s, of which ATOM parts 1.09 s and host parts 0.14 s; CPU 0.79 s. Not tuned: the runtime uploads every input, weights included, per call. - `.research/check_sensitive.py`: clean. ## Compatibility and limits - RBLN only. furiosa / cuda / cpu behaviour is unchanged. - The `rsub` rewrite and host placement apply only to the RBLN backend. - `check` prints extra fp32 fields; `--reference none` skips them. - The default RBLN placement keeps the layer forward on the host. That is a numerics workaround, not a performance choice. - Not measured: quality against A100 / RNGD, and ATOM performance (per-call weight upload). - Not tested: rebel-compiler 0.12 (needs vendor-index credentials). - Vendor bug-report drafts are at the end of the record: `neg` + `add.Scalar` compile failure, and the rounding bias. They are for the lead to send; nothing was sent. ## Related Issue Closes #15 (`check` passes on 2 layers with the default placement, and a LoRA run on ATOM reduces the loss). The compiler and rounding issues are documented with standalone repros. ## Checklist - [x] One purpose; unrelated formatting, dependency or API changes are in separate PRs - [x] Required checks pass locally (new CPU tests) and on ATOM (`check` pass, LoRA runs) - [x] No IPs, internal hostnames, tokens or site-specific paths (`.research/check_sensitive.py`: clean) - [x] Docs updated (README backend table and limitations, experiment record) - [x] AI assistance stated: implemented with Claude Code; reviewed and run by @gspark-etri 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Background
#42 fixed the cause of #40, a race on furiosa-torch's process-global DRAM_SHARED flag. That PR showed 25-step runs only. #40's done-when asks for a 200-step run of
train --full --resident-layers 8 --backend furiosa:auto,furiosa:autowith no OOM. #34 still needed the 1 card × 4 resident timing on an uncontended host.Changes
experiments/2026-10-22-resident-200-steps.md: record of both runs, with losses, step times and the step-model check.experiments/data/resident-200/: step logs (r200.jsonl,t1c.jsonl), pre-submit checks and results (run200.log), and host load.docs/design.md: adds the low-load 1-card row to the step-model table and notes the 200-step result.STATUS.mdand.research/events.jsonlviareport.py.Verification
b3cb03b, which includes Fix resident-AdamW OOM on DRAM_SHARED: serialize device allocations with graph calls (#40) #42 and Qwen2 on RNGD: pass npu-lora check (fp32 q/k + centred k), no negative loss (host target logit) #43.sha256:37a85288…, driver/firmware 2026.4.0.activeDeadlineSecondsandttlSecondsAfterFinished./dev/rngd. The checks logged load 2–21.Run 1: 2 cards × 8 resident layers per card, 200 steps.
Run 2: 1 card, same Job, back to back, load 10–44. Median of steps 6–25:
The model's terms come from the offload run in the same Job.
Compatibility and limits
npu-lora compilein a CPU Job first.furiosa-llm:latestimage, so the first submit hitImagePullBackOff. I deleted that Job at once and re-pulled the image with a pod that holds no card.T_devfrom before Qwen2 on RNGD: pass npu-lora check (fp32 q/k + centred k), no negative loss (host target logit) #43 and is for reference only.experiments/data/dram-shared/race_repro.py).Related Issue
Closes #40. Related to #34.
🤖 Generated with Claude Code