to use
HF_API_TOKEN=<token> sbatch hf-infer.sbatch
then run
HF_API_TOKEN=<token> HOSTNAME=<hostname of infernce server> python gradio-tgl.py
setup env following hf inference server instructions but chance /usr/local to path to conda env instead.
| import numpy as np | |
| import torch | |
| # This is the unfused version of StableAdamW. It is slower than the fused version (coming). | |
| class StableAdamWUnfused(torch.optim.Optimizer): | |
| def __init__(self, params, lr=0.002, weight_decay=0.2, betas=(0.9, 0.99), eps=1e-6, clip_thresh=1., precision='amp_bfloat16', custom_scalar=65536): | |
| beta1, beta2 = betas[0], betas[1] | |
| defaults = dict(lr=lr, weight_decay=weight_decay, beta1=beta1, beta2=beta2) |