fuse_projections: view a group already in one storage instead of copying it
fuse_projections promises "one weight copy", and on a freshly loaded model it keeps that promise. It torch.cats each group (q/k/v/gate, w1/w3) into a new tensor, the separate linears are dropped, and their storage is freed.
The promise breaks when a group's linears are already slices of one merged tensor that something else keeps, for example code that merged the same linears for its own matmul. Then the copy can't free anything, and building SystemOne holds the language model's projections twice.
This change stacks a group as a view when its weights already sit back to back in one storage, in order, and copies otherwise.
- Freshly loaded models are unchanged. Each linear owns its own storage, so every group still takes the copy path.
- Pre-merged weights get one copy. The fused weight is that same storage, so the model holds the group once.
Repro. Merge each language-model group into one tensor that a separate list keeps, then point the linears at its slices:
merged = torch.cat([lin.weight.detach() for lin in linears])
held.append(merged)
for lin, part in zip(linears, merged.split([lin.out_features for lin in linears])):
lin.weight.data = part
Measured on an RTX 4090 (torch 2.11, transformers 5.19). GPU memory allocated after model.engine is built, in GiB:
Shipped lfm3.py |
This change | |
|---|---|---|
| d1-2B, as loaded | 3.57 | 3.57 |
| d1-2B, groups pre-merged | 5.11 | 3.57 |
| d1-Omni-2B, as loaded | 4.05 | 4.05 |
| d1-Omni-2B, groups pre-merged | 5.59 | 4.05 |
- Size of the duplicate: exactly the 35 language-model groups' weights (1.54 GiB).
- Answers: on the Decision Index's ContractNLI (123 contracts × 17 questions, through
system_one), all 2,091 answers and their probabilities are bit-identical to the shipped file as loaded, both as loaded and pre-merged. - Peak memory: 6.47 → 4.93 GiB on d1-2B, 6.95 → 5.41 on d1-Omni-2B. Latency is unchanged.
- State dict:
state_dict()after the fuse still has the native keys and values (max difference 0), and every fused weight shares the merged tensor (35 of 35). - Compiled: with
model.compile()(freezing), d1-2B pre-merged holds 5.12 GiB with the shipped file and 3.58 with this change, the same as a fresh load. Answers are bit-identical here too.
The vision tower's groups are untouched by the repro, and keep the copy path.
The same change is open on LiquidAI/d1-Omni-2B, whose lfm3.py has the same fuse_projections. If you'd rather make this explicit, an option to skip the fuse would also work.
Thanks for this, a neat catch! It's merged here and in our source, so the next release keeps it.