From a6fbfc880c3de9b57e341db374907e2fedda9fa6 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Mon, 16 Jun 2025 10:42:32 -0700 Subject: [PATCH] gguf: fix write order (#11068) * ggml: test write gguf order * ggml: fix write tensor order --- fs/ggml/gguf.go | 12 ++--- fs/ggml/gguf_test.go | 110 +++++++++++++++++++++++++------------------ 2 files changed, 68 insertions(+), 54 deletions(-) diff --git a/fs/ggml/gguf.go b/fs/ggml/gguf.go index 8e75625e..33b596cc 100644 --- a/fs/ggml/gguf.go +++ b/fs/ggml/gguf.go @@ -527,23 +527,17 @@ func WriteGGUF(f *os.File, kv KV, ts []*Tensor) error { return err } - keys := slices.Collect(maps.Keys(kv)) - slices.Sort(keys) - - for _, key := range keys { + for _, key := range slices.Sorted(maps.Keys(kv)) { if err := ggufWriteKV(f, key, kv[key]); err != nil { return err } } slices.SortStableFunc(ts, func(a, b *Tensor) int { - if i, j := a.block(), b.block(); i < 0 && j > 0 { - return 1 - } else if i > 0 && j < 0 { - return -1 - } else { + if i, j := a.block(), b.block(); i > 0 && j > 0 { return cmp.Compare(i, j) } + return cmp.Compare(a.Name, b.Name) }) var s uint64 diff --git a/fs/ggml/gguf_test.go b/fs/ggml/gguf_test.go index 0e071800..bf767918 100644 --- a/fs/ggml/gguf_test.go +++ b/fs/ggml/gguf_test.go @@ -2,62 +2,82 @@ package ggml import ( "bytes" + "math/rand/v2" "os" - "slices" + "strings" "testing" "github.com/google/go-cmp/cmp" ) func TestWriteGGUF(t *testing.T) { - w, err := os.CreateTemp(t.TempDir(), "*.bin") - if err != nil { - t.Fatal(err) - } - defer w.Close() + r := rand.New(rand.NewPCG(0, 0)) + for range 8 { + t.Run("shuffle", func(t *testing.T) { + t.Parallel() - if err := WriteGGUF(w, KV{ - "general.alignment": uint32(16), - }, []*Tensor{ - {Name: "test.0", Shape: []uint64{2, 3}, WriterTo: bytes.NewBuffer(slices.Repeat([]byte{0}, 2*3*4))}, - {Name: "test.1", Shape: []uint64{2, 3}, WriterTo: bytes.NewBuffer(slices.Repeat([]byte{0}, 2*3*4))}, - {Name: "test.2", Shape: []uint64{2, 3}, WriterTo: bytes.NewBuffer(slices.Repeat([]byte{0}, 2*3*4))}, - {Name: "test.3", Shape: []uint64{2, 3}, WriterTo: bytes.NewBuffer(slices.Repeat([]byte{0}, 2*3*4))}, - {Name: "test.4", Shape: []uint64{2, 3}, WriterTo: bytes.NewBuffer(slices.Repeat([]byte{0}, 2*3*4))}, - {Name: "test.5", Shape: []uint64{2, 3}, WriterTo: bytes.NewBuffer(slices.Repeat([]byte{0}, 2*3*4))}, - }); err != nil { - t.Fatal(err) - } + ts := []*Tensor{ + {Name: "token_embd.weight", Shape: []uint64{2, 3}, WriterTo: bytes.NewBuffer(make([]byte, 2*3))}, + {Name: "blk.0.attn_norm.weight", Shape: []uint64{2, 3}, WriterTo: bytes.NewBuffer(make([]byte, 2*3))}, + {Name: "blk.1.attn_norm.weight", Shape: []uint64{2, 3}, WriterTo: bytes.NewBuffer(make([]byte, 2*3))}, + {Name: "blk.2.attn_norm.weight", Shape: []uint64{2, 3}, WriterTo: bytes.NewBuffer(make([]byte, 2*3))}, + {Name: "blk.3.attn_norm.weight", Shape: []uint64{2, 3}, WriterTo: bytes.NewBuffer(make([]byte, 2*3))}, + {Name: "blk.4.attn_norm.weight", Shape: []uint64{2, 3}, WriterTo: bytes.NewBuffer(make([]byte, 2*3))}, + {Name: "blk.5.attn_norm.weight", Shape: []uint64{2, 3}, WriterTo: bytes.NewBuffer(make([]byte, 2*3))}, + {Name: "output_norm.weight", Shape: []uint64{3, 2}, WriterTo: bytes.NewBuffer(make([]byte, 3*2))}, + {Name: "output.weight", Shape: []uint64{3, 2}, WriterTo: bytes.NewBuffer(make([]byte, 3*2))}, + } - r, err := os.Open(w.Name()) - if err != nil { - t.Fatal(err) - } - defer r.Close() + r.Shuffle(len(ts), func(i, j int) { + ts[i], ts[j] = ts[j], ts[i] + }) - ff, err := Decode(r, 0) - if err != nil { - t.Fatal(err) - } + w, err := os.CreateTemp(t.TempDir(), strings.ReplaceAll(t.Name(), "/", "_")+"*.bin") + if err != nil { + t.Fatal(err) + } + defer w.Close() - if diff := cmp.Diff(ff.KV(), KV{ - "general.alignment": uint32(16), - "general.parameter_count": uint64(36), - }); diff != "" { - t.Errorf("Mismatch (-want +got):\n%s", diff) - } + if err := WriteGGUF(w, KV{ + "general.alignment": uint32(16), + }, ts); err != nil { + t.Fatal(err) + } - if diff := cmp.Diff(ff.Tensors(), Tensors{ - Offset: 336, - items: []*Tensor{ - {Name: "test.0", Offset: 0, Shape: []uint64{2, 3}}, - {Name: "test.1", Offset: 32, Shape: []uint64{2, 3}}, - {Name: "test.2", Offset: 64, Shape: []uint64{2, 3}}, - {Name: "test.3", Offset: 96, Shape: []uint64{2, 3}}, - {Name: "test.4", Offset: 128, Shape: []uint64{2, 3}}, - {Name: "test.5", Offset: 160, Shape: []uint64{2, 3}}, - }, - }, cmp.AllowUnexported(Tensors{})); diff != "" { - t.Errorf("Mismatch (-want +got):\n%s", diff) + r, err := os.Open(w.Name()) + if err != nil { + t.Fatal(err) + } + defer r.Close() + + ff, err := Decode(r, 0) + if err != nil { + t.Fatal(err) + } + + if diff := cmp.Diff(KV{ + "general.alignment": uint32(16), + "general.parameter_count": uint64(54), + }, ff.KV()); diff != "" { + t.Errorf("Mismatch (-want +got):\n%s", diff) + } + + if diff := cmp.Diff(Tensors{ + Offset: 608, + items: []*Tensor{ + {Name: "blk.0.attn_norm.weight", Offset: 0, Shape: []uint64{2, 3}}, + {Name: "blk.1.attn_norm.weight", Offset: 32, Shape: []uint64{2, 3}}, + {Name: "blk.2.attn_norm.weight", Offset: 64, Shape: []uint64{2, 3}}, + {Name: "blk.3.attn_norm.weight", Offset: 96, Shape: []uint64{2, 3}}, + {Name: "blk.4.attn_norm.weight", Offset: 128, Shape: []uint64{2, 3}}, + {Name: "blk.5.attn_norm.weight", Offset: 160, Shape: []uint64{2, 3}}, + {Name: "output.weight", Offset: 192, Shape: []uint64{3, 2}}, + {Name: "output_norm.weight", Offset: 224, Shape: []uint64{3, 2}}, + {Name: "token_embd.weight", Offset: 256, Shape: []uint64{2, 3}}, + }, + }, ff.Tensors(), cmp.AllowUnexported(Tensors{})); diff != "" { + t.Errorf("Mismatch (-want +got):\n%s", diff) + } + }) } }