Apple Metal: Mamba parameter dtypes

github.com/openfluke/welvet/nativemetalverified Apple M5

Why it exists

Exercise native storage formats without changing the activation interface.

What it is

Run a seeded Mamba forward pass across all 34 native parameter dtypes, or select one with -dtype.

Storage versus arithmetic

-dtype all sweeps native parameter representations. It does not claim native Float64 GPU arithmetic or integer activation tensors: the Metal execution interface uses Float32 activations. Quantized FormatNone dtype storage and packed quantization formats are different coverage claims.

go run . -dtype all
go run . -dtype float16

Requirements and library installation: chapter 82.

Go example

examples/83-apple-mamba-dtypes/main.go

Run:cd welvet/examples/83-apple-mamba-dtypes && source ../env.sh && go run .

// Float32 tensors with independently selected native parameter encodings.
package main

import (
	"flag"
	"fmt"
	"github.com/openfluke/welvet/core"
	"github.com/openfluke/welvet/layers/mamba"
	"github.com/openfluke/welvet/nativemetal"
	"github.com/openfluke/welvet/quant"
	"log"
	"runtime"
	"strings"
)

func main() {
	if e := run(); e != nil {
		log.Fatal(e)
	}
}
func run() error {
	selected := flag.String("dtype", "all", "native parameter dtype name, or all 34")
	flag.Parse()
	if runtime.GOOS != "darwin" || runtime.GOARCH != "arm64" {
		fmt.Println("SKIP: Apple Metal needs Apple silicon macOS and cgo")
		return nil
	}
	count := 0
	for _, dt := range core.AllDTypes {
		if *selected != "all" && !strings.EqualFold(*selected, dt.String()) {
			continue
		}
		if e := forward(dt); e != nil {
			return fmt.Errorf("%s: %w", dt, e)
		}
		count++
	}
	if count == 0 {
		return fmt.Errorf("unknown dtype %q", *selected)
	}
	fmt.Printf("Mamba dtypes=%d; activations=float32; parameters=FormatNone native dtype\n", count)
	return nil
}
func forward(dt core.DType) error {
	const d, inner = 8, 16
	in, out := make([]float32, 2*inner*d), make([]float32, d*inner)
	for i := range in {
		in[i] = float32(i%17-8) * .01
	}
	for i := range out {
		out[i] = float32(i%13-6) * .01
	}
	model, e := mamba.NewConfigured(mamba.Config{DModel: d, DState: 4, Expand: 2, SeqLen: 3}, dt, quant.FormatNone, in, out, nil, nil)
	if e != nil {
		return e
	}
	s, e := nativemetal.Mount(model)
	if e != nil {
		return e
	}
	defer s.Close()
	x := core.NewTensor[float32](1, 3, d)
	for i := range x.Data {
		x.Data[i] = float32(i%7-3) * .05
	}
	post, e := s.Forward(x)
	if e != nil {
		return e
	}
	var checksum float64
	for _, v := range post.Data {
		checksum += float64(v)
	}
	fmt.Printf("%-10s shape=%v checksum=%.8g\n", dt, post.Shape, checksum)
	return nil
}

Output

exit 0 · last run via go run .

float64    shape=[1 3 8] checksum=-0.016678239
float32    shape=[1 3 8] checksum=-0.016678239
float16    shape=[1 3 8] checksum=-0.016669162
bfloat16   shape=[1 3 8] checksum=-0.016553075
fp8e4m3    shape=[1 3 8] checksum=-0.017004712
fp8e5m2    shape=[1 3 8] checksum=-0.016270325
int64      shape=[1 3 8] checksum=-0.016678241
int32      shape=[1 3 8] checksum=-0.016678239
int16      shape=[1 3 8] checksum=-0.016677922
int8       shape=[1 3 8] checksum=-0.016622871
uint64     shape=[1 3 8] checksum=-0.016678238
uint32     shape=[1 3 8] checksum=-0.01667824
uint16     shape=[1 3 8] checksum=-0.01667843
uint8      shape=[1 3 8] checksum=-0.016727349
int4       shape=[1 3 8] checksum=-0.01510581
uint4      shape=[1 3 8] checksum=-0.017548652
fp4        shape=[1 3 8] checksum=0
int2       shape=[1 3 8] checksum=-0.032299955
uint2      shape=[1 3 8] checksum=-0.015402747
ternary    shape=[1 3 8] checksum=-0.02989793
binary     shape=[1 3 8] checksum=-0.0070436111
int        shape=[1 3 8] checksum=-0.016678241
uint       shape=[1 3 8] checksum=-0.016678238
uintptr    shape=[1 3 8] checksum=-0.016678238
complex64  shape=[1 3 8] checksum=-0.016678239
complex128 shape=[1 3 8] checksum=-0.016678239
nf4        shape=[1 3 8] checksum=-0.016108489
fp6        shape=[1 3 8] checksum=-0.01633804
int6       shape=[1 3 8] checksum=-0.01633804
uint6      shape=[1 3 8] checksum=-0.016878628
int5       shape=[1 3 8] checksum=-0.016140702
uint5      shape=[1 3 8] checksum=-0.017025853
int3       shape=[1 3 8] checksum=-0.018383351
uint3      shape=[1 3 8] checksum=-0.017474993
Mamba dtypes=34; activations=float32; parameters=FormatNone native dtype

Validate (harness)

cd welvet/examples/83-apple-mamba-dtypes && source ../env.sh && go run .