import json, time, tinker
from common import *
print(tinker.AdamParams(learning_rate=1e-4))
sc = service(); val, train, test = splits(); order = train_order(train, 0)
d = [mk(order[0])]
t0 = time.time()
tc = retry(lambda: sc.create_lora_training_client(base_model=MODEL, rank=RANK, **opt_kw("dimuon")), "create")
print("create", time.time() - t0, flush=True)
print("info", tc.get_info(), flush=True)
p = tc.save_state("smoke-init").result(); print("state", p, flush=True)
fb = tc.forward_backward(d, loss_fn="cross_entropy").result(); print("fb", fb.metrics, time.time() - t0, flush=True)
r = tc.optim_step(step_params("dimuon", 1e-3)).result(); print("optim", r.metrics, flush=True)
p2 = tc.save_state("smoke-step1").result(); print("state1", p2, flush=True)
w = tc.save_weights_for_sampler(name="smoke-s1").result(); print("sampler", w, time.time() - t0, flush=True)