fix(lora): sync base weights on the distributed path when the weight checker is armed - #2520
Draft
Arist12 wants to merge 1 commit into
Draft
fix(lora): sync base weights on the distributed path when the weight checker is armed#2520Arist12 wants to merge 1 commit into
Arist12 wants to merge 1 commit into
Conversation
…checker is armed The distributed updater skips base-weight sync for LoRA unconditionally, but the weight checker scrambles the engine's weights first and then asserts the sync restored them, so the scrambled values stay in place and every base tensor fails with a mean absolute error around 0.5. update_weight_from_tensor already carries this carve-out; the distributed path does not. No test covered disaggregated LoRA, so this adds one alongside the colocated case.
Arist12
force-pushed
the
fix/lora-disaggregated-base-weight-check
branch
from
August 13, 2026 16:31
b0ce3ca to
5a97fd4
Compare
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
1 participant
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.
Problem
DistBucketedWeightUpdateMixinskips base-weight sync for LoRA unconditionally, on the grounds that the rollout engines already loaded the frozen base from--hf-checkpoint.--check-weight-update-equalworks by scrambling the engine's weights and then asserting the sync restored them. With the base skipped, the scrambled values stay in place and every base tensor fails with a mean absolute error around 0.5.update_weight_from_tensoralready carries the carve-out for this; the distributed path does not.Reproduced with
examples/lora/run-qwen2.5-3B-megatron-lora-disaggregated.shplus--ci-test.Fix
Mirror the colocated sibling, including its
skip_base_syncspelling (update_weight_from_tensor.py), so the two updaters read the same.Test
tests/e2e/lora/test_lora_qwen2.5_0.5B_disaggregated.py— no test covered disaggregated LoRA at all.One thing worth deciding rather than assuming
This makes the checker path send base weights that production LoRA never sends. The colocated file's own
# TODO: implement lora weight checkersuggests the real answer is a checker that compares adapters only. This PR restores parity between the two updaters; happy to take the adapter-only checker instead if that is the preferred direction.