[gspark-etri]

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

• One purpose; unrelated formatting, dependency or API changes are in separate PRs
• Required checks pass locally (new CPU tests) and on ATOM (check pass, LoRA runs)
• No IPs, internal hostnames, tokens or site-specific paths (.research/check_sensitive.py: clean)
• Docs updated (README backend table and limitations, experiment record)
• AI assistance stated: implemented with Claude Code; reviewed and run by @gspark-etri

🤖 Generated with Claude Code