AMD ROCm: weight return and ENTITY save

github.com/openfluke/welvet/nativerocmlinux/windows amd64

Why it exists

Return trained weights to Go before executing or saving the Go model.

What it is

Train natively, explicitly SyncToGo, save an ENTITY, reload it and compare the native parameter snapshots.

Weight handoff

session, err := nativerocm.Mount(layer)
// Check err; defer session.Close().
// Repeat session.TrainMSE(...) as needed.
err = session.SyncToGo()
// Check err, then save through Go ENTITY IO.

Mount copies model state into a native session. The original Go parameters stay unchanged while that explicit session trains. SyncToGo() returns parameter bytes, scales and Float64 biases for Go saving. This saves model parameters, not a portable checkpoint of live Line, proxy, KV or rule state.

Use -out /path/model.entity to choose the output location. The default uses a temporary directory and prints it.

Requirements: chapter 96.

Go example

examples/99-amd-save-weights/main.go

Run:cd welvet/examples/99-amd-save-weights && source ../env.sh && go run .

// Explicit mount ownership, weight return to Go, ENTITY save and reload.
package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"reflect"
	"runtime"

	"github.com/openfluke/welvet/core"
	"github.com/openfluke/welvet/layers/dense"
	"github.com/openfluke/welvet/layers/parallel"
	"github.com/openfluke/welvet/model/entity"
	"github.com/openfluke/welvet/nativerocm"
	"github.com/openfluke/welvet/quant"
)

func main() {
	if e := run(); e != nil {
		log.Fatal(e)
	}
}

func run() error {
	out := flag.String("out", "", "ENTITY destination; default is a new temporary directory")
	flag.Parse()
	if !rocmPlatformOK() {
		fmt.Println("SKIP: AMD ROCm needs linux/amd64 or windows/amd64 and cgo")
		return nil
	}
	l, e := dense.NewConfigured(2, 2, core.ActivationLinear, core.DTypeFloat32, quant.FormatNone, []float32{.3, -.2, .1, .4})
	if e != nil {
		return e
	}
	model, e := parallel.NewStack(l)
	if e != nil {
		return e
	}
	before, e := l.Weights.FlattenF32()
	if e != nil {
		return e
	}
	s, e := nativerocm.Mount(model)
	if e != nil {
		fmt.Printf("SKIP: ROCm unavailable (%v)\n", e)
		return nil
	}
	defer s.Close()
	x := &core.Tensor[float32]{Shape: []int{1, 2}, Data: []float32{.5, .25}}
	target := core.NewTensor[float32](1, 2)
	for i := 0; i < 5; i++ {
		if _, e = s.TrainMSE(x, target, parallel.ModeNormalBP, .1); e != nil {
			return e
		}
	}
	unchanged, e := l.Weights.FlattenF32()
	if e != nil {
		return e
	}
	if !reflect.DeepEqual(before, unchanged) {
		return fmt.Errorf("explicit mount unexpectedly changed Go parameters before SyncToGo")
	}
	fmt.Println("Go parameters unchanged before SyncToGo")
	if e = s.SyncToGo(); e != nil {
		return e
	}
	after, e := l.Weights.FlattenF32()
	if e != nil {
		return e
	}
	if reflect.DeepEqual(before, after) {
		return fmt.Errorf("expected trained parameters after SyncToGo")
	}
	if *out == "" {
		dir, e := os.MkdirTemp("", "welvet-rocm-entity-")
		if e != nil {
			return e
		}
		*out = filepath.Join(dir, "trained.entity")
	}
	if e = entity.WriteCameralFile(*out, model, parallel.ModeNormalBP); e != nil {
		return e
	}
	loaded, mode, e := entity.LoadCameral(*out)
	if e != nil {
		return e
	}
	if mode != parallel.ModeNormalBP {
		return fmt.Errorf("training mode changed during save")
	}
	restored, e := nativerocm.Mount(loaded)
	if e != nil {
		return e
	}
	defer restored.Close()
	want, e := s.Weights()
	if e != nil {
		return e
	}
	got, e := restored.Weights()
	if e != nil {
		return e
	}
	if !reflect.DeepEqual(want, got) {
		return fmt.Errorf("saved parameters changed on reload")
	}
	fmt.Printf("weights returned to Go; saved=%s; reload=exact\n", *out)
	return nil
}

func rocmPlatformOK() bool {
	return runtime.GOARCH == "amd64" && (runtime.GOOS == "linux" || runtime.GOOS == "windows")
}

Output

exit 0 · last run via go run .

SKIP: ROCm unavailable (nativerocm: load: libamdhip64.so.7: cannot open shared object file: No such file or directory)

Validate (harness)

cd welvet/examples/99-amd-save-weights && source ../env.sh && go run .