Birch-san

Install latest Nvidia driver -- 525 is CUDA 12-era, but is backwards-compatible with 11.8:

apt install nvidia-driver-525 nvidia-dkms-525

Install CUDA 11.8 (use the runfile install, tell it not to replace your driver if you have a newer driver, switch to gcc 11):
https://developer.nvidia.com/cuda-11-8-0-download-archive

Install cuDNN 8 (requires registration):

I started with a basic Ubuntu 22.10 installation, where I chose in the installer to use ZFS as my volume manager.
I wanted to encrypt my home folder.

I followed the article (and comments, including Christoph Hagemann's) from:
https://talldanestale.dk/2020/04/06/zfs-and-homedir-encryption/

To achieve:

	from torch import FloatTensor, LongTensor, Tensor, Size, lerp, zeros_like
	from torch.linalg import norm

	# adapted to PyTorch from:
	# https://gist.github.com/dvschultz/3af50c40df002da3b751efab1daddf2c
	# most of the extra complexity is to support:
	# - many-dimensional vectors
	# - v0 or v1 with last dim all zeroes, or v0 ~colinear with v1
	# - falls back to lerp()
	# - conditional logic implemented with parallelism rather than Python loops

	# suppose we want to matmul some matrix against some column vector..
	# the usual way is this:
	mat = np.array([[1.3e2, 8.e2], [1.6e1, 5.e-1]])
	vec = np.array([2.1e2, 3.6e-4])
	mat @ vec
	array([27300.288 , 3360.00018])

	# but what if we wanted to exploit some efficiencies..
	# - addition can use less cycles / energy / silicon / time than multiplication
	# - for machine learning training: we want to represent a wide range of exponents, but don't need such range on the mantissa

	from torch import Tensor, FloatTensor
	from typing import Protocol, Optional
	from k_diffusion.external import CompVisDenoiser
	from k_diffusion.sampling import sample_heun

	class DiffusionModel(Protocol):
	def __call__(self, x: Tensor, sigma: Tensor, **kwargs) -> Tensor: ...

	class DiffusionModelMixin(DiffusionModel):
	inner_model: DiffusionModel

	import torch
	from torch import einsum, matmul, bmm
	import time

	repeats = 10

	batch_duration = 0
	for ix in range(repeats):
	attn = torch.rand(16, 4096, 4096, dtype=torch.float, device="mps")
	v = torch.rand(16, 4096, 40, dtype=torch.float, device="mps")

	import torch
	from torch import einsum, tensor, matmul, bmm, baddbmm, empty
	import time

	scale=2
	repeats = 10

	# both einsum 0s use the same plan, so whichever batch runs first has to pay the price of warmup
	# uncomment this to run a warmup before either batch runs, for fairer comparison of batch avg time
	# q = torch.rand(16, 4096, 40, dtype=torch.float, device="mps")

	Screenshots in comments below
	https://canary.discord.com/channels/930499730843250783/950914873800396841/1026450454068084798

	import matplotlib.pyplot as plt

	# …

	latents: FloatTensor = self.inner_model(x, sigma, cond=cond, **kwargs)
	unscaled: Tensor = latents / self.scale_factor
	chs = [torch.histogram(c) for c in unscaled[0].flatten(1)]
	h = torch.histogram(unscaled[0].ravel())
	plt.figure(figsize=(10,2))
	plt.title('Per-channel latent values after denoising sigma %.3f at CFG scale %d' % (sigma.item(), cfg_scale))

	CFG20
	sigma: 14.615 absmax: 6.942 std: 2.517 min: -6.840 max: 6.942 shape: [1, 4, 64, 64]
	sigma: 8.560 absmax: 16.164 std: 3.112 min: -16.164 max: 12.644 shape: [1, 4, 64, 64]
	sigma: 8.560 absmax: 28.488 std: 4.065 min: -19.071 max: 28.488 shape: [1, 4, 64, 64]
	sigma: 4.797 absmax: 12.463 std: 1.584 min: -10.655 max: 12.463 shape: [1, 4, 64, 64]
	sigma: 4.797 absmax: 24.024 std: 1.559 min: -24.024 max: 11.797 shape: [1, 4, 64, 64]
	sigma: 2.551 absmax: 5.675 std: 1.151 min: -5.675 max: 4.937 shape: [1, 4, 64, 64]
	sigma: 2.551 absmax: 10.035 std: 1.201 min: -10.035 max: 5.773 shape: [1, 4, 64, 64]
	sigma: 1.274 absmax: 4.725 std: 1.095 min: -3.790 max: 4.725 shape: [1, 4, 64, 64]
	sigma: 1.274 absmax: 4.997 std: 1.104 min: -4.997 max: 4.111 shape: [1, 4, 64, 64]