File size: 4,391 Bytes
e484d86 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | import torch, sys, os, time, json, random
sys.path.insert(0, "/workspace/fractus-cte")
os.chdir("/workspace/fractus-cte")
from fractus.continuous_engine import ContinuousThoughtEngine
import torch.nn.functional as F
GPU = int(os.environ.get("GPU_ID", "0"))
LB_COEF = float(os.environ.get("LB_COEF", "0.02"))
GATE_TEMP = float(os.environ.get("GATE_TEMP", "2.5"))
LR = float(os.environ.get("LR", "5e-4"))
EMA_BETA = 0.98
SS_PROB = float(os.environ.get("SS_PROB", "0.2"))
SS_RATE = float(os.environ.get("SS_RATE", "0.3"))
torch.manual_seed(42 + GPU)
device = torch.device("cuda:0")
TARGET = dict(d_model=1280, n_heads=20, d_head=64, n_levels=2, n_oscillators=16,
coupling_rank=8, n_experts=128, top_k=2, expert_d_ff=2048, siren_rank=64, n_layers=16)
print(f"GPU {GPU}: SS sequential LR={LR} SS_PROB={SS_PROB}", flush=True)
ck = torch.load(f"checkpoints/fractus_1b_gpu{GPU}.pt", map_location="cpu", weights_only=False)
sd = ck.get("model_state", ck)
clean = {(k[10:] if k.startswith("_orig_mod.") else k): v for k,v in sd.items()}
eng = ContinuousThoughtEngine(vocab_size=50257, **{k: TARGET[k] for k in TARGET})
own = eng.state_dict()
for k,v in clean.items():
if k in own and own[k].shape == v.shape:
own[k] = v
elif k in own and v.dim()>=1 and own[k].dim()>=1 and v.shape[0]>own[k].shape[0] and v.shape[1:]==own[k].shape[1:]:
own[k] = v[:own[k].shape[0]].contiguous()
eng.load_state_dict(own, strict=False)
with torch.no_grad():
for blk in eng.blocks:
blk.moe.temperature = GATE_TEMP
eng = eng.to(device)
eng.reset_thought(2)
opt = torch.optim.SGD(eng.parameters(), lr=LR, momentum=0.9)
tokens = torch.load(f"data/shard_gpu{GPU}.pt", weights_only=False).to(torch.int64)
B, seq_len = 2, 128
step_tokens = B * seq_len
man = json.load(open("/workspace/RESUME_MANIFEST_SS.json"))
start_token = int(man["gpus"][str(GPU)]["start_token"])
print(f"GPU {GPU}: RESUME {start_token}", flush=True)
t0=time.time(); ema_tf=None; ema_ss=None; n=0; tok_sess=0
for start in range(start_token, len(tokens)-step_tokens-seq_len-1, step_tokens):
chunk = tokens[start:start+step_tokens].view(B, seq_len).to(device)
target = tokens[start+1:start+step_tokens+1].view(B, seq_len).to(device)
with torch.autocast("cuda", dtype=torch.bfloat16):
out = eng.tick_chunk_train(chunk)
logits, lb = out if isinstance(out, tuple) else (out, eng.last_lb_loss)
ce_tf = F.cross_entropy(logits.reshape(-1, logits.size(-1)), target.reshape(-1))
loss = ce_tf + LB_COEF * lb
opt.zero_grad(); loss.backward()
torch.nn.utils.clip_grad_norm_(eng.parameters(), 1.0); opt.step()
tf_v=float(ce_tf.item()); lb_v=float(lb.detach().item()) if torch.is_tensor(lb) else float(lb)
ema_tf = tf_v if ema_tf is None else EMA_BETA*ema_tf+(1-EMA_BETA)*tf_v
ce_ss_v=None
if random.random() < SS_RATE:
with torch.no_grad():
samp = torch.multinomial(torch.softmax(logits.detach().float().reshape(-1, logits.size(-1))/0.9, -1), 1).view(B, seq_len)
mixed = chunk.clone()
use_ss = torch.rand(B, seq_len, device=device) < SS_PROB
use_ss[:,0]=False
prev = torch.cat([chunk[:,:1], samp[:,:-1]], 1)
mixed = torch.where(use_ss, prev, mixed)
with torch.autocast("cuda", dtype=torch.bfloat16):
out2 = eng.tick_chunk_train(mixed)
logits2, lb2 = out2 if isinstance(out2, tuple) else (out2, eng.last_lb_loss)
ce_ss = F.cross_entropy(logits2.reshape(-1, logits2.size(-1)), target.reshape(-1))
loss2 = 0.5*ce_ss + LB_COEF*lb2
opt.zero_grad(); loss2.backward()
torch.nn.utils.clip_grad_norm_(eng.parameters(), 1.0); opt.step()
ce_ss_v=float(ce_ss.item())
ema_ss = ce_ss_v if ema_ss is None else EMA_BETA*ema_ss+(1-EMA_BETA)*ce_ss_v
n+=1; tok_sess+=step_tokens
if n%50==0:
tps=tok_sess/max(time.time()-t0,1e-6)
extra=f" ss={ce_ss_v:.3f} ema_ss={ema_ss:.3f}" if ce_ss_v is not None and ema_ss is not None else ""
print(f"GPU {GPU}: {start+step_tokens:>12,} tf={tf_v:.3f} ema_tf={ema_tf:.3f}{extra} lb={lb_v:.3f} {tps:.0f} tok/s [ss]", flush=True)
if n%1000==0:
torch.save({"model_state": eng.state_dict(), "config": {**TARGET, "gpu": GPU, "ss": True}}, f"checkpoints/fractus_1b_gpu{GPU}.pt")
print(f"GPU {GPU}: saved [ss]", flush=True)
print("DONE", flush=True)
|