fuse_projections: view a group already in one storage instead of copying it (#3)

- fuse_projections: view a group already in one storage instead of copying it (33b2fffb16d2e24874a7dafee0c1d4b16fe975ab)


Co-authored-by: Chao-Chun (Joe) Hsu <joe32140@users.noreply.huggingface.co>

lfm3.py

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