J3ff/IQ.Pilot
1
0
Fork 0
forked from IQ.Lvbs/IQ.Pilot
Code Pull Requests Activity

IQ.Pilot Release Commit @ bec7652

Browse Source
This commit is contained in:
IQ.Lvbs history cleanup
2026-08-22 23:42:41 -05:00
commit 58039e647c
4603 changed files with 1236178 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
artifacts/package_sources/tinygrad/test/external/mlperf_stable_diffusion/external_test_eval.py vendored Normal file
View File

@@ -0,0 +1,58 @@
import unittest, os
import numpy as np
from pathlib import Path
from tempfile import TemporaryDirectory
from tinygrad import Device, Tensor
from tinygrad.helpers import getenv, Context
from tinygrad.nn.state import safe_save, torch_load, get_parameters
from examples.mlperf.model_eval import eval_stable_diffusion, vae_decode
from examples.stable_diffusion import AutoencoderKL
def set_eval_params():
# override these as needed from cli
for k,v in {"MODEL": "stable_diffusion", "GPUS": "8", "EVAL_SAMPLES": "600", "CONTEXT_BS": "816", "DENOISE_BS": "600", "DECODE_BS": "384",
"INCEPTION_BS": "560", "CLIP_BS": "240", "DATADIR": "/raid/datasets/stable_diffusion", "CKPTDIR": "/raid/weights/stable_diffusion"}.items():
os.environ[k] = getenv(k, v)
class TestEval(unittest.TestCase):
def test_eval_ckpt(self):
set_eval_params()
with TemporaryDirectory(prefix="test-eval") as tmp:
os.environ["EVAL_CKPT_DIR"] = tmp
# NOTE Although this checkpoint has the original fully trained model from StabilityAI, we are using mlperf code that uses different
# GroupNorm num_groups. Therefore, eval results may not reflect eval results on the original model.
# The purpose of using this checkpoint is to have reproducible eval outputs.
# Eval code expects file and weight names in a specific format, as .safetensors (not .ckpt), which is why we resave the checkpoint
sd_v2 = torch_load(Path(getenv("CKPTDIR", "")) / "sd" / "512-base-ema.ckpt")["state_dict"]
sd_v2 = {k.replace("model.diffusion_model.", "", 1): v for k,v in sd_v2.items() if k.startswith("model.diffusion_model.")}
safe_save(sd_v2, f"{tmp}/0.safetensors")
clip, fid, ckpt = eval_stable_diffusion()
assert ckpt == 0
if Device.DEFAULT == "NULL":
assert clip == 0
assert fid > 0 and fid < 1000
else:
# observed:
# clip=0.08369670808315277, fid=301.05236173709545 (if SEED=12345, commit=c01b2c93076e80ae6d1ebca64bb8e83a54dadba6)
# clip=0.08415728807449341, fid=300.3710877072948 (if SEED=12345, commit=179c7fcfe132f1a6344b57c9d8cef4eded586867)
# clip=0.0828116238117218, fid=301.241909555543 (if SEED=98765, commit=c01b2c93076e80ae6d1ebca64bb8e83a54dadba6)
np.testing.assert_allclose(fid, 301.147, rtol=0.1, atol=0)
np.testing.assert_allclose(clip, 0.08325, rtol=0.1, atol=0)
# only tested on 8xMI300x system
@unittest.skipUnless(getenv("HANG_OK"), "expected to hang")
def test_decoder_beam_hang(self):
set_eval_params()
for k,v in {"BEAM": "2", "HCQDEV_WAIT_TIMEOUT_MS": "300000", "BEAM_UOPS_MAX": "8000", "BEAM_UPCAST_MAX": "256", "BEAM_LOCAL_MAX": "1024",
"BEAM_MIN_PROGRESS": "5", "IGNORE_JIT_FIRST_BEAM": "1"}.items():
os.environ[k] = getenv(k, v)
with Context(BEAM=int(os.environ["BEAM"])): # necessary because helpers.py has already set BEAM=0 and cached getenv for "BEAM"
GPUS = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 8))]
vae = AutoencoderKL()
for p in get_parameters(vae): p.to_(GPUS).realize()
x = Tensor.zeros(48,4,64,64).contiguous().to(GPUS).realize()
x.uop = x.uop.unshard(0)
for _ in range(2): vae_decode(x, vae)
if __name__=="__main__":
unittest.main()

114
artifacts/package_sources/tinygrad/test/external/mlperf_stable_diffusion/external_test_models.py vendored Normal file
View File

@@ -0,0 +1,114 @@
import unittest
import numpy as np
from pathlib import Path
from tinygrad import Tensor, dtypes, Device
from tinygrad.helpers import getenv
from tinygrad.nn.state import get_parameters
from extra.models import clip
from examples.mlperf.initializers import gelu_erf, init_stable_diffusion, attn_f32_softmax
from typing import Literal
clip_params = {"dims": 1024, "n_heads": 16, "layers": 24, "return_pooled": False, "ln_penultimate": True}
def get_cond_stage_model(GPUS:list[str]|None=None) -> clip.FrozenOpenClipEmbedder:
clip.gelu = gelu_erf
model = clip.FrozenOpenClipEmbedder(**clip_params)
if GPUS and len(GPUS) > 1:
for p in get_parameters(model): p.to_(GPUS)
return model
def get_tokens(BS:int) -> Tensor: return Tensor([0] * 77 * BS, dtype=dtypes.int32).reshape(-1, 77)
class TestOpenClip(unittest.TestCase):
def test_tokenizer(self):
prompt = "Beautiful is better than ugly.\nExplicit is better than implicit.\nSimple is better than complex.\nComplex is better than complicated."
model = get_cond_stage_model()
tokens = model.tokenizer.encode(prompt, pad_with_zeros=True)
expected = [49406, 1215, 533, 1539, 1126, 36070, 269, 33228, 533, 1539, 1126, 15269, 529, 269, 4129, 533, 1539, 1126, 16099, 269, 6324, 533,
1539, 1126, 11460, 14589, 269, 49407] + [0]*49
self.assertEqual(tokens, expected)
def test_clip_gelu_init(self):
for resblock in get_cond_stage_model().model.transformer.resblocks:
self.assertEqual(resblock.mlp.gelu, gelu_erf)
def test_multigpu_clip_embed(self):
BS = 304
GPUS = [f"{Device.DEFAULT}:{i}" for i in range(8)]
model = get_cond_stage_model(GPUS)
tokens = get_tokens(BS)
embeds = model.embed_tokens(tokens.shard(GPUS, axis=0)).realize()
self.assertEqual(embeds.shape, (BS, 77, 1024))
self.assertEqual(embeds.dtype, dtypes.float32)
def test_multigpu_clip_score(self):
BS = 240
GPUS = [f"{Device.DEFAULT}:{i}" for i in range(8)]
vision_cfg = {'width': 1280, 'layers': 32, 'd_head': 80, 'image_size': 224, 'patch_size': 14}
text_cfg = {'width': 1024, 'n_heads': 16, 'layers': 24, 'vocab_size': 49408, 'ctx_length': 77}
clip.gelu = gelu_erf
clip_encoder = clip.OpenClipEncoder(1024, text_cfg, vision_cfg)
for p in get_parameters(clip_encoder): p.to_(GPUS)
tokens = get_tokens(BS)
imgs = Tensor.zeros(BS,3,224,224).contiguous()
scores = clip_encoder.get_clip_score(tokens.shard(GPUS, axis=0), imgs.shard(GPUS, axis=0)).realize()
self.assertEqual(scores.shape, (BS,))
self.assertEqual(scores.dtype, dtypes.float32)
class TestInitStableDiffusion(unittest.TestCase):
def setUp(self):
# NOTE: set env variable based on where checkpoints are on the system
self.CKPTDIR = Path(getenv("CKPTDIR", "/raid/weights/stable_diffusion"))
def helper_test_init(self, version:Literal["v2-mlperf-train", "v2-mlperf-eval"]):
model, unet, sqrt_acp, sqrt_omacp = init_stable_diffusion(version, self.CKPTDIR / "sd" / "512-base-ema.ckpt", ["CPU"])
with self.subTest("test that StableDiffusion has correct models"):
self.assertEqual(model.model.diffusion_model, unet)
has_encoder = True if version=="v2-mlperf-eval" else False
self.assertEqual(hasattr(model, "first_stage_model"), has_encoder, "only the eval model uses the encoder")
self.assertTrue(isinstance(model.cond_stage_model, clip.FrozenOpenClipEmbedder))
with self.subTest("test for mlperf unique attributes"):
self.assertEqual(model.cond_stage_model.tokenizer.version, 'sd_mlperf_v5_0')
self.assertEqual(unet.out[0].num_groups, 16)
self.assertEqual(unet.input_blocks[1][1].norm.eps, 1e-6)
self.assertEqual(unet.input_blocks[1][1].transformer_blocks[0].attn1.attn, attn_f32_softmax)
with self.subTest("test loaded clip parameters"):
sample = model.cond_stage_model.model.transformer.resblocks[8].mlp.c_fc.bias.flatten()[42:46].numpy()
expected = np.array([-0.49812260270118713, -0.3039605915546417, -0.40284937620162964, -0.45069342851638794], dtype=np.float32)
np.testing.assert_allclose(sample, expected, rtol=1e-7, atol=0, err_msg="loaded clip parameters are incorrect")
if version=="v2-mlperf-train":
with self.subTest("test that zero_module worked"):
self.assertTrue((unet.out[2].weight == 0).all().item(), "expected all zeroes")
self.assertTrue((unet.out[2].bias == 0).all().item(), "expected all zeroes")
elif version=="v2-mlperf-eval":
with self.subTest("test loaded vae parameters"):
sample = model.first_stage_model.decoder.up[0]['block'][1].conv2.weight.flatten()[42:46].numpy()
expected = np.array([0.08192943036556244, 0.040095631033182144, 0.07541035860776901, 0.1475081741809845], dtype=np.float32)
np.testing.assert_allclose(sample, expected, rtol=1e-7, atol=0, err_msg="loaded vae parameters are incorrect")
with self.subTest("check schedules"):
expected = np.array([0.9995748996734619, 0.06826484948396683], dtype=np.float32)
np.testing.assert_allclose(sqrt_acp[[0,-1]].numpy(), expected, rtol=1e-7, atol=0, err_msg="sqrt_acp is incorrect")
expected = np.array([0.029155133292078972, 0.9976672530174255], dtype=np.float32)
np.testing.assert_allclose(sqrt_omacp[[0,-1]].numpy(), expected, rtol=1e-7, atol=0, err_msg="sqrt_omacp is incorrect")
with self.subTest("check mixed precision"):
out = unet.input_blocks[2][1].proj_in(Tensor.randn(320, dtype=dtypes.float32))
self.assertEqual(out.dtype, dtypes.bfloat16, "expected float32 to be downcast to bfloat16 by Linear")
out = unet.out[2](Tensor.randn(304,320,64,64, dtype=dtypes.float32))
self.assertEqual(out.dtype, dtypes.bfloat16, "expected float32 to be downcast to bfloat16 by Conv2d")
out = unet.input_blocks[1][1].transformer_blocks[0].norm1(Tensor.randn(320, dtype=dtypes.bfloat16))
self.assertEqual(out.dtype, dtypes.float32, "expected bfloat16 to be upcast to float32 by LayerNorm")
out = unet.input_blocks[5][0].in_layers[0](Tensor.randn(304, 640, dtype=dtypes.bfloat16))
self.assertEqual(out.dtype, dtypes.float32, "expected bfloat16 to be upcast to float32 by GroupNorm")
def test_train_model(self):
self.helper_test_init("v2-mlperf-train")
def test_eval_model(self):
self.helper_test_init("v2-mlperf-eval")
if __name__=="__main__":
unittest.main()

23
artifacts/package_sources/tinygrad/test/external/mlperf_stable_diffusion/external_test_train.py vendored Normal file
View File

@@ -0,0 +1,23 @@
import unittest, os
from tempfile import TemporaryDirectory
from tinygrad import Context
from tinygrad.helpers import getenv
from examples.mlperf.model_train import train_stable_diffusion
class TestTrain(unittest.TestCase):
def test_train_to_ckpt(self):
# train for num_steps, save checkpoint, and stop training
num_steps = 42
os.environ.update({"MODEL": "stable_diffusion", "TOTAL_CKPTS": "1", "CKPT_STEP_INTERVAL": str(num_steps), "GPUS": "8", "BS": "304"})
# NOTE: update these based on where data/checkpoints are on your system
if not getenv("DATADIR", ""): os.environ["DATADIR"] = "/raid/datasets/stable_diffusion"
if not getenv("CKPTDIR", ""): os.environ["CKPTDIR"] = "/raid/weights/stable_diffusion"
with TemporaryDirectory(prefix="test-train") as tmp:
os.environ["UNET_CKPTDIR"] = tmp
with Context(TRAINING=1):
saved_ckpts = train_stable_diffusion()
expected_ckpt = f"{tmp}/{num_steps}.safetensors"
assert len(saved_ckpts) == 1 and saved_ckpts[0] == expected_ckpt
if __name__=="__main__":
unittest.main()