Repository navigation
Commit b8b23fd
fix(peft)!: make expert LoRA exact under suppressed tensor-parallel communication
Complete the adapter arithmetic that explicit_expert_comm suppression leaves
undone, using autograd-visible collectives on the groups each adapter already
carries: reduce_from(copy_to(x)) over expert-TP for the row-parallel-input
case (all-reduce on both passes), copy_to(x) for the column-parallel case
(backward-only all-reduce restoring the suppressed dL/dz sum), and a
differentiable last-dim all-gather (backward reduce-scatter-sum) replacing
the grouped adapter's autograd-invisible bare all_gather in all three
backends. On the output side, keep linear_out unsharded-gather-free and
zero-embed the local shard at its rank offset, so the downstream sum (the
dispatcher over expert-TP; SharedExpertMLP.post_forward_comm over dense TP)
counts the delta exactly once with no scaling factor; MultiLoRALinear gets
the same treatment on suppressed-comm row-parallel bases.
Where no correct completion exists without new machinery, refuse loudly
instead of training wrong: DoRA on expert or suppressed-comm TP-sharded
linears, SharedOuterGroupedExpertAdapter at ETP>1, expert-adapter dropout at
ETP>1 (per-rank RNG on dispatcher-replicated tokens), and LoRALinear.weight
on grouped adapter classes.
BREAKING CHANGE: the four configurations above now raise NotImplementedError
where they previously ran with silently wrong numerics (or an accidental
AttributeError). Checkpoint format, weight layout, and export paths are
unchanged.
The strict-xfail markers from the previous commit are deleted here; every
demonstration now asserts for real. Verified on 2x H100 (--runxfail,
TF32 off): CPU suite 13/13, distributed 11/11 on both ranks, refusal and
wiring tests 26 passed, full peft regression 419 passed. At ETP=1 every new
branch short-circuits at group size 1; peak memory is unchanged by
construction (the embed buffer replaces the removed gather buffer).
The fix design originates from Mohammad Taufeeque's earlier implementation
of the same defect class.
Co-authored-by: Taufeeque <9taufeeque9@gmail.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Tigist Diriba <tigist@far.ai>1 parent e0b1a27 commit b8b23fd
10 files changed
Lines changed: 410 additions & 122 deletions
File tree
- src/megatron/bridge/peft
- tests/unit_tests/peft
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
21 | 21 | | |
22 | 22 | | |
23 | 23 | | |
24 | | - | |
| 24 | + | |
25 | 25 | | |
26 | 26 | | |
27 | 27 | | |
| |||
90 | 90 | | |
91 | 91 | | |
92 | 92 | | |
| 93 | + | |
| 94 | + | |
| 95 | + | |
| 96 | + | |
| 97 | + | |
| 98 | + | |
| 99 | + | |
93 | 100 | | |
| 101 | + | |
| 102 | + | |
| 103 | + | |
| 104 | + | |
| 105 | + | |
| 106 | + | |
| 107 | + | |
| 108 | + | |
94 | 109 | | |
95 | 110 | | |
96 | 111 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
35 | 35 | | |
36 | 36 | | |
37 | 37 | | |
| 38 | + | |
| 39 | + | |
| 40 | + | |
| 41 | + | |
| 42 | + | |
| 43 | + | |
| 44 | + | |
| 45 | + | |
| 46 | + | |
| 47 | + | |
| 48 | + | |
| 49 | + | |
| 50 | + | |
38 | 51 | | |
39 | 52 | | |
40 | 53 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
115 | 115 | | |
116 | 116 | | |
117 | 117 | | |
| 118 | + | |
| 119 | + | |
| 120 | + | |
| 121 | + | |
| 122 | + | |
| 123 | + | |
| 124 | + | |
| 125 | + | |
| 126 | + | |
| 127 | + | |
| 128 | + | |
| 129 | + | |
118 | 130 | | |
119 | 131 | | |
120 | 132 | | |
| |||
155 | 167 | | |
156 | 168 | | |
157 | 169 | | |
| 170 | + | |
| 171 | + | |
| 172 | + | |
| 173 | + | |
| 174 | + | |
| 175 | + | |
| 176 | + | |
| 177 | + | |
| 178 | + | |
| 179 | + | |
| 180 | + | |
| 181 | + | |
| 182 | + | |
| 183 | + | |
| 184 | + | |
158 | 185 | | |
159 | 186 | | |
160 | 187 | | |
| |||
226 | 253 | | |
227 | 254 | | |
228 | 255 | | |
| 256 | + | |
| 257 | + | |
| 258 | + | |
229 | 259 | | |
230 | 260 | | |
| 261 | + | |
| 262 | + | |
231 | 263 | | |
232 | 264 | | |
233 | 265 | | |
| |||
0 commit comments