Commit 974f073
committed
Fix -sot: mirror output_norm to split context for all devices
Ikawrakow's review: output norm tensors remain on buft_output.buft
(the per-device offload buft) instead of being on the split context.
With -sot, the output weight IS split across devices, but the norm
weights stay on a single device and never get mirrored.
Fix: in get_context_for_tensor(), when split_output_tensor is on,
route output_norm (and output_norm_bias, output.bias) to
ctx_output_split (= split_buft context). This puts them in
split_tensors so prepare_split_tensors(-1, ...) at tensor-load
time creates per-device mirrored copies. build_output() already
handles output_norm->extra correctly: line 2107 checks for it
and line 2115 uses per-device norm copies when available.1 parent 3a3e736 commit 974f073
1 file changed
Lines changed: 3 additions & 0 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
447 | 447 | | |
448 | 448 | | |
449 | 449 | | |
| 450 | + | |
| 451 | + | |
| 452 | + | |
450 | 453 | | |
451 | 454 | | |
452 | 455 | | |
| |||
0 commit comments