Update251203 (#233)
This commit is contained in:
carrot
@@ -758,6 +758,27 @@ def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0
|
||||
batch.append(tokens)
|
||||
yield Tensor.stack(batch, dim=0)
|
||||
|
||||
def batch_load_llama3_small(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
|
||||
if val:
|
||||
dataset = BlendedGPTDataset([
|
||||
base_dir / "c4-validation-91205-samples.en_text_document",
|
||||
], [
|
||||
1.0
|
||||
], samples, seqlen, seed, False)
|
||||
else:
|
||||
dataset = BlendedGPTDataset([
|
||||
base_dir / "c4-train.en_6_text_document",
|
||||
], [
|
||||
1.0
|
||||
], samples, seqlen, seed, True)
|
||||
|
||||
for b in range(math.ceil(samples / bs)):
|
||||
batch = []
|
||||
for i in range(bs):
|
||||
tokens = dataset.get(b * bs + i)
|
||||
batch.append(tokens)
|
||||
yield Tensor.stack(batch, dim=0)
|
||||
|
||||
if __name__ == "__main__":
|
||||
def load_unet3d(val):
|
||||
assert not val, "validation set is not supported due to different sizes on inputs"
|
||||
|
||||
@@ -243,31 +243,49 @@ def eval_mrcnn():
|
||||
|
||||
def eval_llama3():
|
||||
from extra.models.llama import Transformer
|
||||
from examples.llama3 import MODEL_PARAMS
|
||||
from examples.llama3 import MODEL_PARAMS, load, convert_from_huggingface
|
||||
from tinygrad.helpers import tqdm
|
||||
|
||||
bs = 4
|
||||
sequence_length = 512
|
||||
BASEDIR = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
|
||||
BS = getenv("BS", 4)
|
||||
SMALL = getenv("SMALL", 0)
|
||||
SEQLEN = getenv("SEQLEN", 8192)
|
||||
MODEL_PATH = Path(getenv("MODEL_PATH", "/raid/weights/llama31_8b/"))
|
||||
|
||||
model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
|
||||
params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
|
||||
params = params | {"vocab_size": 32000} if not SMALL else params
|
||||
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
|
||||
model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
|
||||
|
||||
# load weights
|
||||
weights = load(str(MODEL_PATH / "model.safetensors.index.json"))
|
||||
if "model.embed_tokens.weight" in weights:
|
||||
print("converting from huggingface format")
|
||||
weights = convert_from_huggingface(weights, params["n_layers"], params["n_heads"], params["n_kv_heads"])
|
||||
|
||||
load_state_dict(model, weights, strict=False, consume=True)
|
||||
|
||||
@TinyJit
|
||||
def eval_step(model, tokens):
|
||||
logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
|
||||
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
|
||||
return loss.flatten()
|
||||
return loss.flatten().float()
|
||||
|
||||
from examples.mlperf.dataloader import batch_load_llama3
|
||||
iter = batch_load_llama3(bs, 5760, sequence_length, Path(getenv("BASEDIR", "/raid/datasets/c4/")), True)
|
||||
if SMALL:
|
||||
from examples.mlperf.dataloader import batch_load_llama3_small
|
||||
iter = batch_load_llama3_small(BS, 5760, SEQLEN, BASEDIR, val=True)
|
||||
else:
|
||||
from examples.mlperf.dataloader import batch_load_llama3
|
||||
iter = batch_load_llama3(BS, 5760, SEQLEN, BASEDIR, val=True)
|
||||
|
||||
losses = []
|
||||
for tokens in tqdm(iter, total=5760//bs):
|
||||
for tokens in tqdm(iter, total=5760//BS):
|
||||
GlobalCounters.reset()
|
||||
losses += eval_step(model, tokens).tolist()
|
||||
tqdm.write(f"loss: {np.mean(losses)}")
|
||||
|
||||
log_perplexity = Tensor(losses).mean()
|
||||
print(f"Log Perplexity: {log_perplexity.item()}")
|
||||
log_perplexity = np.mean(losses)
|
||||
print(f"Log Perplexity: {log_perplexity}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
# inference only
|
||||
|
||||
@@ -4,7 +4,7 @@ import multiprocessing
|
||||
|
||||
from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
|
||||
from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW, Profiling
|
||||
from tinygrad.nn.state import get_parameters, get_state_dict, safe_load, safe_save
|
||||
from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
|
||||
from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
|
||||
|
||||
from extra.lr_scheduler import LRSchedulerGroup
|
||||
@@ -252,6 +252,10 @@ def train_resnet():
|
||||
print(f"epoch global_ops: {steps_in_train_epoch * GlobalCounters.global_ops:_}, "
|
||||
f"epoch global_mem: {steps_in_train_epoch * GlobalCounters.global_mem:_}")
|
||||
# if we are doing beam search, run the first eval too
|
||||
if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
|
||||
min_time = min(step_times)
|
||||
assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
|
||||
|
||||
if (TRAIN_BEAM or EVAL_BEAM) and e == start_epoch: break
|
||||
return
|
||||
if MLLOGGER and RUNMLPERF:
|
||||
@@ -344,6 +348,8 @@ def train_resnet():
|
||||
print(f"saving ckpt to {fn}")
|
||||
safe_save(get_training_state(model, optimizer_group, scheduler_group), fn)
|
||||
|
||||
|
||||
|
||||
def train_retinanet():
|
||||
from contextlib import redirect_stdout
|
||||
from examples.mlperf.dataloader import batch_load_retinanet
|
||||
@@ -1290,12 +1296,14 @@ def train_llama3():
|
||||
from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
|
||||
|
||||
config = {}
|
||||
BASEDIR = config["BASEDIR"] = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
|
||||
BS = config["BS"] = getenv("BS", 16)
|
||||
grad_acc = config["GRADIENT_ACC_STEPS"] = getenv("GRADIENT_ACC_STEPS", 1)
|
||||
GBS = config["GLOBAL_BATCH_SIZE"] = BS * grad_acc
|
||||
SEED = config["SEED"] = getenv("SEED", 5760)
|
||||
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
|
||||
TRAIN_ON_VAL = config["TRAIN_ON_VAL"] = getenv("TRAIN_ON_VAL", 0)
|
||||
SMALL = config["SMALL"] = getenv("SMALL", 0)
|
||||
SAMPLES = config["SAMPLES"] = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
|
||||
EVAL_FREQ = config["EVAL_FREQ"] = getenv("EVAL_FREQ", 46080)
|
||||
EVAL_BS = config["EVAL_BS"] = getenv("EVAL_BS", 16)
|
||||
@@ -1311,13 +1319,14 @@ def train_llama3():
|
||||
|
||||
opt_gradient_clip_norm = 1.0
|
||||
opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
|
||||
opt_learning_rate_decay_steps = getenv("DECAY_STEPS", math.ceil(1_200_000 * 1152 / GBS) - opt_learning_rate_warmup_steps)
|
||||
opt_learning_rate_decay_steps = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS)) - opt_learning_rate_warmup_steps
|
||||
opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152) # NOTE: cannot change for benchmark
|
||||
opt_end_learning_rate = 8e-7
|
||||
opt_end_learning_rate = getenv("END_LR", 8e-7)
|
||||
|
||||
# TODO: confirm weights are in bf16
|
||||
# vocab_size from the mixtral tokenizer
|
||||
params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}
|
||||
params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
|
||||
params = params | {"vocab_size": 32000} if not SMALL else params
|
||||
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
|
||||
model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
|
||||
|
||||
@@ -1353,6 +1362,15 @@ def train_llama3():
|
||||
b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
|
||||
scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
|
||||
|
||||
if resume_ckpt := getenv("RESUME_CKPT"):
|
||||
fn = f"./ckpts/llama3_{resume_ckpt}.safe"
|
||||
print(f"loading initial checkpoint from {fn}")
|
||||
load_state_dict(model, safe_load(fn), realize=False)
|
||||
|
||||
fn = f"./ckpts/llama3_{resume_ckpt}_optim.safe"
|
||||
print(f"loading optim checkpoint from {fn}")
|
||||
load_state_dict(scheduler, safe_load(fn), realize=False)
|
||||
|
||||
@TinyJit
|
||||
@Tensor.train()
|
||||
def train_step(model, tokens:Tensor, grad_acc:int):
|
||||
@@ -1403,43 +1421,55 @@ def train_llama3():
|
||||
# ** data iters **
|
||||
def fake_data(bs, samples):
|
||||
for _ in range(samples // bs):
|
||||
yield Tensor.randint(bs, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
|
||||
yield Tensor.randint(bs, SEQLEN + 1, low=0, high=params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
|
||||
|
||||
def get_train_iter():
|
||||
if getenv("FAKEDATA", 0):
|
||||
return fake_data(GBS, SAMPLES)
|
||||
else:
|
||||
from examples.mlperf.dataloader import batch_load_llama3
|
||||
return batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))
|
||||
if SMALL:
|
||||
from examples.mlperf.dataloader import batch_load_llama3_small
|
||||
return batch_load_llama3_small(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
|
||||
else:
|
||||
from examples.mlperf.dataloader import batch_load_llama3
|
||||
return batch_load_llama3(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
|
||||
|
||||
def get_eval_iter():
|
||||
if getenv("FAKEDATA", 0):
|
||||
return fake_data(EVAL_BS, 5760)
|
||||
else:
|
||||
from examples.mlperf.dataloader import batch_load_llama3
|
||||
return batch_load_llama3(EVAL_BS, 5760, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=True)
|
||||
if SMALL:
|
||||
from examples.mlperf.dataloader import batch_load_llama3_small
|
||||
return batch_load_llama3_small(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
|
||||
else:
|
||||
from examples.mlperf.dataloader import batch_load_llama3
|
||||
return batch_load_llama3(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
|
||||
|
||||
iter = get_train_iter()
|
||||
i, sequences_seen = 0, 0
|
||||
i, sequences_seen = resume_ckpt, 0
|
||||
for tokens in tqdm(iter, total=SAMPLES//GBS):
|
||||
t = time.perf_counter()
|
||||
GlobalCounters.reset()
|
||||
loss, lr = train_step(model, tokens, grad_acc)
|
||||
loss = loss.float().item()
|
||||
# above as tqdm.write f-string
|
||||
|
||||
i += 1
|
||||
sequences_seen += tokens.shape[0]
|
||||
|
||||
tqdm.write(f"{loss:.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
|
||||
if (fname:=getenv("LOSS_FILE", "")):
|
||||
with open(fname, "a") as f:
|
||||
f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
|
||||
|
||||
if getenv("CKPT") and (i % 200 == 0 or i == 10):
|
||||
if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
|
||||
tqdm.write("saving checkpoint")
|
||||
if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
|
||||
fn = f"{ckpt_dir}/llama3_{i}.safe"
|
||||
safe_save(get_state_dict(model), fn)
|
||||
|
||||
i += 1
|
||||
sequences_seen += tokens.shape[0]
|
||||
tqdm.write("saving optim checkpoint")
|
||||
fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
|
||||
safe_save(get_state_dict(scheduler), fn)
|
||||
|
||||
if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
|
||||
tqdm.write(f"evaluating after {sequences_seen} sequences")
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
#!/usr/bin/env bash
|
||||
# adapted from https://github.com/mlcommons/training/blob/4bdf5c8ed218ad76565a2ba1ac27c919ccc6d689/stable_diffusion/README.md
|
||||
|
||||
# setup dirs
|
||||
|
||||
DATA=/raid/datasets/stable_diffusion
|
||||
|
||||
LAION=$DATA/laion-400m/webdataset-moments-filtered
|
||||
COCO=$DATA/coco2014
|
||||
mkdir -p $LAION $COCO
|
||||
|
||||
CKPT=/raid/weights/stable_diffusion
|
||||
mkdir -p $CKPT/clip $CKPT/sd $CKPT/inception
|
||||
|
||||
# download data
|
||||
|
||||
# if rclone isn't installed system-wide / in your PATH, put the executable path in quotes below
|
||||
#RCLONE=""
|
||||
RCLONE="rclone"
|
||||
|
||||
## VAE-encoded image latents, from 6.1M image subset of laion-400m
|
||||
## about 1 TB for whole download
|
||||
$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
|
||||
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/ ${LAION} --include="*.tar" -P
|
||||
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/sha512sums.txt ${LAION} -P
|
||||
cd $LAION && grep -E '\.tar$' sha512sums.txt | sha512sum -c --quiet - && \
|
||||
echo "All .tar files verified" || { echo "Checksum failure when validating downloaded Laion moments"; exit 1; }
|
||||
|
||||
## prompts and FID statistics from 30k image subset of coco2014
|
||||
## 33 MB
|
||||
$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
|
||||
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k.tsv ${COCO} -P
|
||||
|
||||
$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
|
||||
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k_stats.npz ${COCO} -P
|
||||
|
||||
# download checkpoints
|
||||
|
||||
## clip (needed for text and vision encoders for validation)
|
||||
CLIP_WEIGHTS_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin"
|
||||
CLIP_WEIGHTS_SHA256="9a78ef8e8c73fd0df621682e7a8e8eb36c6916cb3c16b291a082ecd52ab79cc4"
|
||||
CLIP_CONFIG_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/raw/main/open_clip_config.json"
|
||||
wget -N -P ${CKPT}/clip ${CLIP_WEIGHTS_URL}
|
||||
wget -N -P ${CKPT}/clip ${CLIP_CONFIG_URL}
|
||||
echo "${CLIP_WEIGHTS_SHA256} ${CKPT}/clip/open_clip_pytorch_model.bin" | sha256sum -c
|
||||
|
||||
## sd (needed for latent->image decoder for validation, also has clip text encoder for training)
|
||||
SD_WEIGHTS_URL='https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt'
|
||||
SD_WEIGHTS_SHA256="d635794c1fedfdfa261e065370bea59c651fc9bfa65dc6d67ad29e11869a1824"
|
||||
wget -N -P ${CKPT}/sd ${SD_WEIGHTS_URL}
|
||||
echo "${SD_WEIGHTS_SHA256} ${CKPT}/sd/512-base-ema.ckpt" | sha256sum -c
|
||||
|
||||
## inception (needed for validation)
|
||||
FID_WEIGHTS_URL='https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'
|
||||
FID_WEIGHTS_SHA1="bd836944fd6db519dfd8d924aa457f5b3c8357ff"
|
||||
wget -N -P ${CKPT}/inception ${FID_WEIGHTS_URL}
|
||||
echo "${FID_WEIGHTS_SHA1} ${CKPT}/inception/pt_inception-2015-12-05-6726825d.pth" | sha1sum -c
|
||||
Reference in New Issue
Block a user