diff options
| -rw-r--r-- | gpt2/model.go | 185 |
1 files changed, 123 insertions, 62 deletions
diff --git a/gpt2/model.go b/gpt2/model.go index fed1712..1058ef4 100644 --- a/gpt2/model.go +++ b/gpt2/model.go @@ -4,7 +4,6 @@ import ( "errors" "fmt" _ "llm" - "log" "math" "os" "slices" @@ -21,8 +20,11 @@ const ( ) type Model struct { - name string - deviceID string + name string + deviceID string + session *ort.DynamicAdvancedSession + inputNames []string + outputNames []string } func NewModel(name, deviceID string) *Model { @@ -45,7 +47,43 @@ func (m *Model) SharedLibraryPath() string { func (m *Model) Init() error { ort.SetSharedLibraryPath(m.SharedLibraryPath()) - return ort.InitializeEnvironment() + if err := ort.InitializeEnvironment(); err != nil { + return err + } + + inputNames := make([]string, 0, 3+2*nLayers) + outputNames := make([]string, 0, 1+2*nLayers) + + inputNames = append(inputNames, "input_ids", "position_ids", "attention_mask") + outputNames = append(outputNames, "logits") + + for i := range nLayers { + inputNames = append(inputNames, fmt.Sprintf("past_key_values.%d.key", i), fmt.Sprintf("past_key_values.%d.value", i)) + outputNames = append(outputNames, fmt.Sprintf("present.%d.key", i), fmt.Sprintf("present.%d.value", i)) + } + + m.inputNames = inputNames + m.outputNames = outputNames + + var options *ort.SessionOptions + + if m.deviceID != "" { + if opts, err := SessionsOptionsWithCUDADeviceID(m.deviceID); err != nil { + return err + } else { + options = opts + + defer options.Destroy() + } + } + + if s, err := ort.NewDynamicAdvancedSession(m.name, m.inputNames, m.outputNames, options); err != nil { + return err + } else { + m.session = s + } + + return nil } func (m *Model) Destroy() error { @@ -59,14 +97,18 @@ func (m *Model) Generate(prompt []int64, steps int64, logits *[][]float32) ([]in context := int64(len(prompt)) - cacheNames, cacheValues := emptyCache() + cacheValues := emptyCache() + + defer func() { + destroyValues(cacheValues) + }() token := prompt[0] out := make([]int64, 0, steps+1) for step := range context + steps { - _, _, outputs, err := m.forward(m.name, token, step, cacheNames, cacheValues) + outputs, err := m.forward(token, step, cacheValues) if err != nil { return nil, err @@ -93,91 +135,116 @@ func (m *Model) Generate(prompt []int64, steps int64, logits *[][]float32) ([]in out = append(out, token) } + _ = outputs[0].Destroy() + + destroyValues(cacheValues) + cacheValues = outputs[1:] } return out[:steps], nil } -func (m *Model) forward(model string, token int64, position int64, cacheNames []string, cacheValues []ort.Value) (*ort.Tensor[float32], []string, []ort.Value, error) { - inputNames, inputs, _ := initInputs(token, position) - outputNames, outputs, logits, _ := initOutputs(position) +func (m *Model) forward(token int64, position int64, cacheValues []ort.Value) ([]ort.Value, error) { + var binding *ort.IoBinding + + if b, err := m.session.CreateIoBinding(); err != nil { + return nil, err + } else { + binding = b + + defer binding.Destroy() + } + + var inputs []ort.Value + var outputs []ort.Value + + if in, err := initInputs(token, position); err != nil { + return nil, err + } else { + inputs = in + } + + defer destroyValues(inputs) + + if out, err := initOutputs(position); err != nil { + return nil, err + } else { + outputs = out + } - inputNames = append(inputNames, cacheNames...) inputs = append(inputs, cacheValues...) - var options *ort.SessionOptions + if len(inputs) != len(m.inputNames) { + panic("unexpected input length") + } - if m.deviceID != "" { - if opts, err := SessionsOptionsWithCUDADeviceID(m.deviceID); err != nil { - return nil, nil, nil, err - } else { - options = opts + for i, name := range m.inputNames { + if err := binding.BindInput(name, inputs[i]); err != nil { + return nil, err } } - session, err := ort.NewAdvancedSession( - model, - inputNames, - outputNames, - inputs, - outputs, - options, - ) + var ok bool - if err != nil { - log.Fatal(err) + defer func() { + if !ok { + destroyValues(outputs) + } + }() + + if len(outputs) != len(m.outputNames) { + panic("unexpected output length") } - if options != nil { - if err := options.Destroy(); err != nil { - return nil, nil, nil, err + for i, name := range m.outputNames { + if err := binding.BindOutput(name, outputs[i]); err != nil { + return nil, err } } - defer session.Destroy() - - if err := session.Run(); err != nil { - log.Fatal(err) + if err := m.session.RunWithBinding(binding); err != nil { + return nil, err } - return logits, outputNames, outputs, nil + ok = true + + return outputs, nil +} + +func destroyValues(values []ort.Value) { + for _, v := range values { + _ = v.Destroy() + } } -func emptyCache() ([]string, []ort.Value) { - names := make([]string, 0, 2*nLayers) +func emptyCache() []ort.Value { values := make([]ort.Value, 0, 2*nLayers) shape := []int64{1, int64(nHeads), 0, int64(headDim)} - for i := range nLayers { - kName := fmt.Sprintf("past_key_values.%d.key", i) - vName := fmt.Sprintf("past_key_values.%d.value", i) - + for range nLayers { kTensor, _ := ort.NewEmptyTensor[float32](shape) vTensor, _ := ort.NewEmptyTensor[float32](shape) - names = append(names, kName, vName) values = append(values, ort.Value(kTensor), ort.Value(vTensor)) } - return names, values + return values } -func initInputs(token, position int64) ([]string, []ort.Value, error) { - inputNames := []string{"input_ids", "position_ids", "attention_mask"} - +func initInputs(token, position int64) ([]ort.Value, error) { var tokens *ort.Tensor[int64] var positions *ort.Tensor[int64] var attentionMask *ort.Tensor[int64] if t, err := ort.NewTensor[int64]([]int64{1, 1}, []int64{token}); err != nil { - return nil, nil, err + return nil, err } else { tokens = t } if p, err := ort.NewTensor[int64]([]int64{1, 1}, []int64{position}); err != nil { - return nil, nil, err + return nil, err } else { positions = p } @@ -190,43 +257,37 @@ func initInputs(token, position int64) ([]string, []ort.Value, error) { } if m, err := ort.NewTensor[int64](maskShape, maskData); err != nil { - return nil, nil, err + return nil, err } else { attentionMask = m } inputs := []ort.Value{ort.Value(tokens), ort.Value(positions), ort.Value(attentionMask)} - return inputNames, inputs, nil + return inputs, nil } -func initOutputs(position int64) ([]string, []ort.Value, *ort.Tensor[float32], error) { - outputNames := make([]string, 0, 1+2*nLayers) - outputValues := make([]ort.Value, 0, 1+2*nLayers) +func initOutputs(position int64) ([]ort.Value, error) { + outputs := make([]ort.Value, 0, 1+2*nLayers) logits, err := ort.NewEmptyTensor[float32]([]int64{1, 1, int64(vocabSize)}) if err != nil { - return nil, nil, nil, err + return nil, err } - outputNames = append(outputNames, "logits") - outputValues = append(outputValues, ort.Value(logits)) + outputs = append(outputs, ort.Value(logits)) shape := []int64{1, int64(nHeads), position + 1, int64(headDim)} - for i := range nLayers { - kName := fmt.Sprintf("present.%d.key", i) - vName := fmt.Sprintf("present.%d.value", i) - + for range nLayers { kTensor, _ := ort.NewEmptyTensor[float32](shape) vTensor, _ := ort.NewEmptyTensor[float32](shape) - outputNames = append(outputNames, kName, vName) - outputValues = append(outputValues, ort.Value(kTensor), ort.Value(vTensor)) + outputs = append(outputs, ort.Value(kTensor), ort.Value(vTensor)) } - return outputNames, outputValues, logits, nil + return outputs, nil } func softmax(logits []float32) []float32 { |