diff options
| -rw-r--r-- | go.work | 1 | ||||
| -rw-r--r-- | profile/go.mod | 3 | ||||
| -rw-r--r-- | profile/profile.go | 64 | ||||
| -rw-r--r-- | research/knobloch/cmd/train/main.go | 7 | ||||
| -rw-r--r-- | tokenizer/bpe/cmd/tokenize/main.go | 46 |
5 files changed, 79 insertions, 42 deletions
@@ -7,6 +7,7 @@ use ( ./llmc ./mbpe ./onnx + ./profile ./research/knobloch ./research/lesci ./research/sander diff --git a/profile/go.mod b/profile/go.mod new file mode 100644 index 0000000..3fcd33c --- /dev/null +++ b/profile/go.mod @@ -0,0 +1,3 @@ +module go.jknobloc.com/x/profile + +go 1.25 diff --git a/profile/profile.go b/profile/profile.go new file mode 100644 index 0000000..5654ba1 --- /dev/null +++ b/profile/profile.go @@ -0,0 +1,64 @@ +package profile + +import ( + "flag" + "log" + "os" + "runtime" + "runtime/pprof" +) + +var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to `file`") +var memprofile = flag.String("memprofile", "", "write memory profile to `file`") + +func CPU() func() { + flag.Parse() + + if *cpuprofile == "" { + return nil + } + + var file *os.File + + if f, err := os.Create(*cpuprofile); err != nil { + log.Fatal("could not create CPU profile: ", err) + } else { + file = f + } + + if err := pprof.StartCPUProfile(file); err != nil { + log.Fatal("could not start CPU profile: ", err) + } + + return func() { + pprof.StopCPUProfile() + + if err := file.Close(); err != nil { + log.Fatal(err) + } + } +} + +func Mem() { + flag.Parse() + + if *memprofile == "" { + return + } + + var file *os.File + + if f, err := os.Create(*memprofile); err != nil { + log.Fatal("could not create memory profile: ", err) + } else { + file = f + + defer file.Close() + } + + runtime.GC() + + if err := pprof.Lookup("allocs").WriteTo(file, 0); err != nil { + log.Fatal("could not write memory profile: ", err) + } +} diff --git a/research/knobloch/cmd/train/main.go b/research/knobloch/cmd/train/main.go index a946ce4..d2f49f2 100644 --- a/research/knobloch/cmd/train/main.go +++ b/research/knobloch/cmd/train/main.go @@ -10,14 +10,21 @@ import ( "github.com/jonasknobloch/mbpe" "go.jknobloc.com/x/dataset" + "go.jknobloc.com/x/profile" "go.jknobloc.com/x/shelf" "go.jknobloc.com/x/tokenizer/bpe" "go.jknobloc.com/x/tokenizer/bpe/split" ) func main() { + stop := profile.CPU() + // train() // serialize() + + profile.Mem() + + stop() } func train() { diff --git a/tokenizer/bpe/cmd/tokenize/main.go b/tokenizer/bpe/cmd/tokenize/main.go index e6cab4b..f8c81c1 100644 --- a/tokenizer/bpe/cmd/tokenize/main.go +++ b/tokenizer/bpe/cmd/tokenize/main.go @@ -1,44 +1,20 @@ package main import ( - "flag" "log" - "os" - "runtime" - "runtime/pprof" "sync/atomic" "time" "go.jknobloc.com/x/dataset" "go.jknobloc.com/x/llm" + "go.jknobloc.com/x/profile" "go.jknobloc.com/x/shelf" "go.jknobloc.com/x/tokenizer/bpe" "go.jknobloc.com/x/tui" ) -var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to `file`") -var memprofile = flag.String("memprofile", "", "write memory profile to `file`") - func main() { - flag.Parse() - - if *cpuprofile != "" { - var file *os.File - - if f, err := os.Create(*cpuprofile); err != nil { - log.Fatal("could not create CPU profile: ", err) - } else { - file = f - - defer file.Close() - } - - if err := pprof.StartCPUProfile(file); err != nil { - log.Fatal("could not start CPU profile: ", err) - } - - defer pprof.StopCPUProfile() - } + stop := profile.CPU() reader := data() @@ -66,23 +42,9 @@ func main() { processed.Add(1) } - if *memprofile != "" { - var file *os.File - - if f, err := os.Create(*memprofile); err != nil { - log.Fatal("could not create memory profile: ", err) - } else { - file = f - - defer file.Close() - } + profile.Mem() - runtime.GC() - - if err := pprof.Lookup("allocs").WriteTo(file, 0); err != nil { - log.Fatal("could not write memory profile: ", err) - } - } + stop() } func data() *dataset.ParquetReader { |