joe32140 commited on
Commit
33b2fff
·
verified ·
1 Parent(s): cfe5760

fuse_projections: view a group already in one storage instead of copying it

Browse files

`fuse_projections` promises "one weight copy", and on a freshly loaded model it keeps that promise. It `torch.cat`s 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:

```python
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.

Files changed (1) hide show
  1. lfm3.py +15 -1
lfm3.py CHANGED
@@ -135,12 +135,26 @@ def fuse_projections(model) -> None:
135
  names = SHARED_INPUT.get(type(m).__name__, ())
136
  linears = [getattr(m, n, None) for n in names]
137
  if names and all(type(lin) is torch.nn.Linear and lin.bias is None for lin in linears):
138
- m.fused = Fused(torch.cat([lin.weight for lin in linears]), names, [lin.out_features for lin in linears])
139
  for i, name in enumerate(names):
140
  setattr(m, name, Part(m.fused, i))
141
  m.register_state_dict_post_hook(_unfused)
142
 
143
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
144
  def _unfused(module, state_dict, prefix, local_metadata) -> None:
145
  fused = module.fused
146
  weight = state_dict.pop(prefix + "fused.weight")
 
135
  names = SHARED_INPUT.get(type(m).__name__, ())
136
  linears = [getattr(m, n, None) for n in names]
137
  if names and all(type(lin) is torch.nn.Linear and lin.bias is None for lin in linears):
138
+ m.fused = Fused(_stacked([lin.weight for lin in linears]), names, [lin.out_features for lin in linears])
139
  for i, name in enumerate(names):
140
  setattr(m, name, Part(m.fused, i))
141
  m.register_state_dict_post_hook(_unfused)
142
 
143
 
144
+ def _stacked(weights: list[torch.Tensor]) -> torch.Tensor:
145
+ """`weights` stacked by rows: a view when they already sit back to back in one storage, in this order, else a
146
+ copy. A freshly loaded model's linears each own their storage, so this copies and the old weights are freed.
147
+ Linears that are slices of one merged tensor something else keeps would otherwise be held twice (d1-2B: 1.5 GiB)."""
148
+ first = weights[0].detach()
149
+ storage, offset = first.untyped_storage().data_ptr(), first.storage_offset()
150
+ for w in weights:
151
+ if (w.untyped_storage().data_ptr() != storage or w.storage_offset() != offset or not w.is_contiguous()
152
+ or w.dtype != first.dtype or w.shape[1:] != first.shape[1:]):
153
+ return torch.cat(weights)
154
+ offset += w.numel()
155
+ return first.as_strided((sum(w.shape[0] for w in weights), *first.shape[1:]), first.stride())
156
+
157
+
158
  def _unfused(module, state_dict, prefix, local_metadata) -> None:
159
  fused = module.fused
160
  weight = state_dict.pop(prefix + "fused.weight")