84. Apple Metal: training modes

Part: XII · Apple acceleration (welvet.metal)
Package: nativemetal + parallel + architecture
Status: ok: verified macOS/arm64 on Apple M5

When

Train a mounted model using native Metal credit and optimizer kernels.

Where

From example/welvet/84-apple-training:

go run .
go run . -mode StepBP -steps 8
go run . -mode MeshTweenChain
go run . -mode all

Why

Keep one session alive for Step pipeline state; use an actual Grid for Mesh modes.

What

Runs a two-layer Dense model in NormalBP by default. All 29 update modes are selectable. Mesh modes use a disabled ingress plus two active Grid cells and four clock ticks per call.

Requirements

Apple silicon macOS, Go with cgo, and the installed independent Metal dylib. Set WELVET_METAL_LIBRARY or install it from the Metal checkout:

cd /Users/openfluke/gs/welvet.metal
./scripts/macos_install_go.sh

Use the Go v1.3.0 checkout containing nativemetal. Examples 82-86 can also use go run -tags welvet_metal . to exclude the independent C++ binding. Example 87 needs the normal build to keep both native backends. Unsupported operating systems print SKIP. Missing libraries on a supported host report an error.

These are usage examples with small shapes. For speed measurements, use ../apple_stress_test. The results are numerical CPU/GPU comparisons, not model prediction accuracy.

Backend labels in routed examples: cpu-tiled is Go CPU, native_cpu is welvet.cpp, and metal is the independent Apple GPU backend.

Captured sample

mode=NormalBP schedule=Stack calls=8 loss=6.7954462e-06 weights=synced-to-Go

main.go

Download source ↓

// Native Metal MSE training: persistent Stack Line state or an actual Grid clock.
package main

import (
	"flag"
	"fmt"
	"github.com/openfluke/welvet/architecture"
	"github.com/openfluke/welvet/core"
	"github.com/openfluke/welvet/layers/dense"
	"github.com/openfluke/welvet/layers/parallel"
	"github.com/openfluke/welvet/nativemetal"
	"github.com/openfluke/welvet/quant"
	"log"
	"runtime"
)

func main() {
	if e := run(); e != nil {
		log.Fatal(e)
	}
}
func run() error {
	selected := flag.String("mode", "NormalBP", "training mode name, or all 29")
	steps := flag.Int("steps", 8, "training calls; retain the session for Step pipeline state")
	flag.Parse()
	if *steps < 1 {
		return fmt.Errorf("steps must be positive")
	}
	if runtime.GOOS != "darwin" || runtime.GOARCH != "arm64" {
		fmt.Println("SKIP: Apple Metal needs Apple silicon macOS and cgo")
		return nil
	}
	if *selected == "all" {
		for mode := parallel.ModeNormalBP; mode <= parallel.ModeStepTweenSplitSparse; mode++ {
			if e := train(mode, *steps); e != nil {
				return e
			}
		}
		return nil
	}
	mode, e := parallel.ParseTrainMode(*selected)
	if e != nil {
		return e
	}
	if mode < parallel.ModeNormalBP || mode > parallel.ModeStepTweenSplitSparse {
		return fmt.Errorf("choose one of the 29 update modes")
	}
	return train(mode, *steps)
}
func train(mode parallel.TrainMode, steps int) error {
	var ops []any
	for layer := 0; layer < 2; layer++ {
		w := make([]float32, 16)
		for i := range w {
			w[i] = float32(i%5-2) * .04
		}
		l, e := dense.NewConfigured(4, 4, core.ActivationTanh, core.DTypeFloat32, quant.FormatNone, w)
		if e != nil {
			return e
		}
		ops = append(ops, l)
	}
	stack, e := parallel.NewStack(ops...)
	if e != nil {
		return e
	}
	var model any = stack
	schedule := "Stack"
	if mode.RequiresGrid() {
		g := architecture.NewGrid(3, 1, 1, 1)
		g.Cells[0].Layer.IsDisabled = true
		for i, op := range ops {
			g.Cells[i+1].Layer = op.(*dense.Layer).Core
			g.Cells[i+1].Op = op
		}
		model = g
		schedule = "Grid mesh clock (4 ticks/call)"
	} else if mode.IsLineStep() {
		schedule = "Stack Line pipeline"
	}
	s, e := nativemetal.Mount(model)
	if e != nil {
		return e
	}
	defer s.Close()
	x := core.NewTensor[float32](2, 4)
	for i := range x.Data {
		x.Data[i] = float32(i%5-2) * .1
	}
	target := core.NewTensor[float32](2, 4)
	var loss float64
	for i := 0; i < steps; i++ {
		loss, e = s.TrainMSE(x, target, mode, .01, 4)
		if e != nil {
			return fmt.Errorf("%s: %w", mode, e)
		}
	}
	if e = s.SyncToGo(); e != nil {
		return e
	}
	fmt.Printf("mode=%s schedule=%s calls=%d loss=%.8g weights=synced-to-Go\n", mode, schedule, steps, loss)
	return nil
}