[RFC] Make native LoRA adapter sync the default for RL refit and keep merged full-weight sync opt-in
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Summary
For LoRA training with a generation backend that supports native adapters (initially vLLM), NeMo-RL should synchronize the factorized LoRA adapter by default instead of materializing and transferring merged full weights. The current merged path should remain available as an explicit compatibility/performance opt-in.
This proposal is motivated by a reproducible correctness issue on Nemotron Nano V3.5 LoRA GRPO: materializing `W + scale * B @ A` onto the BF16 base-weight grid substantially changes the intended LoRA update and can produce large trainer/generation logprob divergence. This is a finite-precision representation boundary, not evidence of a generic vLLM loader bug.
The upstream vLLM community is tracking adapter-only RL synchronization in [vllm-project/vllm#48297](https://github.com/vllm-project/vllm/issues/48297), and it is listed in the [2026 Q3 vLLM × RL roadmap](https://github.com/vllm-project/vllm/issues/48314). NeMo-RL still needs a framework-level default, validation contract, and fallback policy.
## Current behavior
The DTensor LoRA refit path currently:
1. skips `.lora_A.weight` and `.lora_B.weight` tensors;
2. computes `W_merged = W + scale * B @ A` for each `LinearLoRA` base weight;
3. casts/materializes that full tensor in the generation policy dtype; and
4. transfers the full state to the generation engine.
Current implementation: [`dtensor_policy_worker_v2.py`](https://github.com/NVIDIA-NeMo/RL/blob/d3974b21a6305d2e77b3134f64f8aa32d4980ba5/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py#L117-L170). The corresponding refit metadata also excludes the A/B factors: [`prepare_refit_info`](https://github.com/NVIDIA-NeMo/RL/blob/d3974b21a6305d2e77b3134f64f8aa32d4980ba5/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py#L1086-L1105).
In exact arithmetic, the trainer and merged generation expressions are equivalent:
```text
trainer: x W^T + scale * (x A^T) B^T
generation: x (W + scale * B @ A)^T
```
They are not equivalent after the merged full weight is materialized on a BF16/FP16 grid.
## Evidence
### 1. Step-20 weight-level BF16 merge audit
For a Nemotron Nano V3.5 LoRA GRPO checkpoint:
| Metric | Result |
|---|---:|
| Nonzero intended LoRA-delta elements | 30,910,978,558 |
| Nonzero delta elements that did not change the BF16 base weight | 27,994,898,219 |
| Rounded-back fraction | 90.5662% |
| Effective BF16 merged-delta L2 / intended-delta L2 | 0.5739 |
| Relative L2 error | 0.9114 |
The rounded-back fraction is an element count, not an energy percentage. The L2 statistics show that the surviving update is also substantially changed in aggregate magnitude and direction.
### 2. Same-checkpoint replay
We scored the same fixed 2,044 token positions from the same checkpoint under three representations:
| Comparison | Generation KL | Mean absolute logprob error | Max absolute logprob error |
|---|---:|---:|---:|
| Trainer factorized → vLLM BF16 merged | 0.040052 | 0.16318 | 4.25870 |
| Trainer factorized → trainer BF16 merged | 0.032176 | 0.14602 | 3.88705 |
| Trainer BF16 merged → vLLM BF16 merged | 0.005094 | 0.05537 | 0.70563 |
Once the trainer is forced onto the same BF16 merged representation as vLLM, the residual generation KL drops by approximately 87.3%. Most of the abnormal mismatch can therefore be reproduced before crossing the trainer-to-vLLM transport/loader boundary.
This replay changes both the final storage representation and the factorized-versus-single-GEMM execution order. The weight-level audit independently demonstrates that BF16 materialization itself causes substantial update distortion.
### 3. Native adapter feasibility
An experimental vLLM-native path transferred BF16 LoRA A/B tensors, created and activated a native vLLM adapter, and attached a `LoRARequest` to generation/scoring requests. A 40-step Nemotron Nano V3.5 run completed with generation KL in `[0.0017, 0.0034]` (mean `0.002335`).
That initial 40-step experiment also included final CUDA stream fences, so it is not by itself a single-variable attribution result. Subsequent fixed-checkpoint audits show that stream synchronization and BF16 merge representation are separate concerns; adapter-only refit remains the path that avoids the demonstrated `BF16(W + BA)` information loss.
## Related work
- vLLM adapter lifecycle RFC: [vllm-project/vllm#48297](https://github.com/vllm-project/vllm/issues/48297)
- vLLM weight-reload correctness RFC: [vllm-project/vllm#48312](https://github.com/vllm-project/vllm/issues/48312)
- vLLM × RL roadmap, LoRA lifecycle section: [vllm-project/vllm#48314](https://github.com/vllm-project/vllm/issues/48314)
- NeMo-RL router correction-bias dtype fix: [NVIDIA-NeMo/RL#3540](https://github.com/NVIDIA-NeMo/RL/pull/3540)
Contributor guide
Assessment
This issue has not been assessed yet.