Skip to content

fix(rlhf): preserve base weights during LoRA rollout sync - #10154

Closed
taking-lying-flat wants to merge 1 commit into
modelscope:mainfrom
taking-lying-flat:fix/bf16-lora-rollout-sync
Closed

taking-lying-flat wants to merge 1 commit into
modelscope:mainfrom
taking-lying-flat:fix/bf16-lora-rollout-sync

Conversation

@taking-lying-flat

Copy link
Copy Markdown
Contributor

HF/DeepSpeed rollout synchronization temporarily merges LoRA into the training model and then subtracts the adapter delta to unmerge. BF16 rounding makes that subtraction an inexact inverse: a synchronization alone can change frozen base weights, even without an optimizer step.

Preserve the original base parameter storage and merge into temporary tensors inside the existing parameter-group gather context. Restore the original parameters and PEFT merge/cache state when the context exits, including on failures. Apply this to both full-weight sync and the existing merge in adapter-only sync.

The change keeps the existing parameter groups, transport calls, adapter dtype alignment, and outgoing payload layout. It does not introduce additional transfer batches or write model copies to disk. Temporary base tensors do add memory proportional to the merged parameters in the current group; with move_model_batches=None, that group can cover the whole model.

Validation:

  • Six local tests using the actual source methods and real PEFT layers passed: 100 repeated syncs preserve the original BF16 base exactly; merge and sender failures restore it; coverage includes DoRA, Embedding, Conv2d, nested parameter wrappers, and tied base parameters.
  • Compared outgoing tensor names, values, dtypes, byte counts, and call counts against the unmodified source for full-weight and adapter-only sync, with and without grouping. All matched.
  • Repository pre-commit checks passed. Network transport and ZeRO gathering were mocked in the local tests; a full distributed DeepSpeed/vLLM run was not performed.

Only production code is included. Related numerical mechanism: #10087 discusses BF16 merge/unmerge perturbation in the Megatron save path; this change addresses the HF rollout synchronization path.

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