Update Mo_jax.py

Mo_jax.py

@@ -89,43 +89,49 @@ def create_batch_iter(inputs, targets, batch_size, rng):
 
 def shard(xs): return xs.reshape(NUM_DEVICES, -1, xs.shape[1])
 
-# ------------------
-# Model
-# ------------------
 class SwiGLU(nn.Module):
     d_model: int
-    dtype: Any = DTYPE
     @nn.compact
-    def __call__(self,x):
-        proj = nn.Dense(self.d_model*2,dtype=self.dtype)(x)
-        x_val, x_gate = jnp.split(proj,2,-1)
+    def __call__(self, x):
+        x_f32 = x.astype(jnp.float32)
+        proj = nn.Dense(self.d_model*2, dtype=jnp.float32)(x_f32)
+        x_val, x_gate = jnp.split(proj, 2, axis=-1)
         out = x_val * nn.silu(x_gate)
-        return nn.Dense(self.d_model,dtype=self.dtype)(out)
+        out = nn.Dense(self.d_model, dtype=jnp.float32)(out)
+        return out.astype(x.dtype)
 
 class LoU(nn.Module):
-    d_model:int
-    dtype:Any=DTYPE
+    d_model: int
+    clip_value: float = 5.0
+    eps: float = 1e-6
     @nn.compact
-    def __call__(self,x):
-        residual = x
-        x_norm = nn.LayerNorm(epsilon=1e-5,dtype=self.dtype)(x)
-        Q=nn.Dense(self.d_model,dtype=self.dtype)
-        K=nn.Dense(self.d_model,dtype=self.dtype)
-        V=nn.Dense(self.d_model,dtype=self.dtype)
-        q,k,v = Q(x_norm),K(x_norm),V(x_norm)
-        g_q = (jnp.tanh(q)+1)/2; g_k=(jnp.tanh(k)+1)/2
-        score = g_q*g_k
-        alpha_dynamic = nn.Dense(1,dtype=self.dtype)(x_norm)
+    def __call__(self, x):
+        x_f32 = x.astype(jnp.float32)
+        residual = x_f32
+        x_norm = nn.LayerNorm(epsilon=1e-5, dtype=jnp.float32)(x_f32)
+        Q = nn.Dense(self.d_model, dtype=jnp.float32)
+        K = nn.Dense(self.d_model, dtype=jnp.float32)
+        V = nn.Dense(self.d_model, dtype=jnp.float32)
+        q,k,v = Q(x_norm), K(x_norm), V(x_norm)
+        g_q = (jnp.tanh(q)+1)/2
+        g_k = (jnp.tanh(k)+1)/2
+        score = g_q * g_k
+        alpha_dynamic = nn.Dense(1, dtype=jnp.float32)(x_norm)
         # EMA scan along seq axis
         score_t = jnp.transpose(score,(1,0,2))
         alpha_t = jnp.transpose(alpha_dynamic,(1,0,2))
-        def step(prev,cur): s,a=cur; new=a*s+(1-a)*prev; return new,new
-        init = score_t[0]; _,ema_seq=jax.lax.scan(step,init,(score_t[1:],alpha_t[1:]))
-        ema_full=jnp.concatenate([init[None,...],ema_seq],0)
+        def step(prev, cur):
+            s, a = cur
+            new = a*s + (1-a)*prev
+            return new,new
+        init = score_t[0]
+        _, ema_seq = jax.lax.scan(step, init, (score_t[1:], alpha_t[1:]))
+        ema_full = jnp.concatenate([init[None,...], ema_seq], 0)
         ema = jnp.transpose(ema_full,(1,0,2))
-        out = v*ema + residual
-        out = nn.LayerNorm(epsilon=1e-5,dtype=self.dtype)(out)
-        return SwiGLU(self.d_model,self.dtype)(out)
+        out = v * ema + residual
+        out = nn.LayerNorm(epsilon=1e-5, dtype=jnp.float32)(out)
+        return SwiGLU(self.d_model)(out).astype(x.dtype)
+
 
 class Lo(nn.Module):
     d_model:int
@@ -161,28 +167,26 @@ class ReLM(nn.Module):
         logits=jnp.einsum("bld,vd->blv",x,self.token_embed.embedding)
         return logits
 
-# ------------------
-# Loss & metrics
-# ------------------
-def smoothed_ce(logits,targets,pad_id,eps=0.1):
-    vocab=logits.shape[-1]
-    logits=logits.reshape(-1,vocab)
-    targets=targets.reshape(-1)
-    mask=(targets!=pad_id).astype(jnp.float32)
-    one_hot=jax.nn.one_hot(targets,vocab)
-    smooth=(1-eps)*one_hot+eps/vocab
-    log_probs=jax.nn.log_softmax(logits)
-    loss=-jnp.sum(smooth*log_probs,axis=-1)*mask
-    return jnp.sum(loss)/(jnp.sum(mask)+1e-8)
+def smoothed_ce(logits, targets, pad_id, eps=0.1):
+    logits = logits.astype(jnp.float32)
+    targets = targets.astype(jnp.int32)
+    vocab = logits.shape[-1]
+    mask = (targets != pad_id).astype(jnp.float32)
+    one_hot = jax.nn.one_hot(targets, vocab)
+    smooth = (1-eps)*one_hot + eps/vocab
+    log_probs = jax.nn.log_softmax(logits, axis=-1)
+    loss = -jnp.sum(smooth * log_probs, axis=-1) * mask
+    return jnp.sum(loss) / (jnp.sum(mask)+1e-8)
 
-def masked_ppl(logits,targets,pad_id,eps=0.1):
-    vocab=logits.shape[-1]
-    logits=logits.reshape(-1,vocab)
-    targets=targets.reshape(-1)
-    mask=(targets!=pad_id).astype(jnp.float32)
-    one_hot=jax.nn.one_hot(targets,vocab)
-    smooth=(1-eps)*one_hot+eps/vocab
-    loss=-jnp.sum(smooth*jax.nn.log_softmax(logits),axis=-1)*mask
+def masked_ppl(logits, targets, pad_id, eps=0.1):
+    logits = logits.astype(jnp.float32)
+    targets = targets.astype(jnp.int32)
+    vocab = logits.shape[-1]
+    mask = (targets != pad_id).astype(jnp.float32)
+    one_hot = jax.nn.one_hot(targets, vocab)
+    smooth = (1-eps)*one_hot + eps/vocab
+    log_probs = jax.nn.log_softmax(logits, axis=-1)
+    loss = -jnp.sum(smooth*log_probs, axis=-1) * mask
     return jnp.exp(jnp.sum(loss)/(jnp.sum(mask)+1e-8))
 
 # ------------------
@@ -264,7 +268,13 @@ for epoch in range(EPOCHS):
 # ------------------
 save_dir="./checkpoints"
 os.makedirs(save_dir,exist_ok=True)
-checkpoints.save_checkpoint(save_dir,jax.tree_map(lambda x:np.array(x),state),step=global_step,keep=3)
+# 기존
+# checkpoints.save_checkpoint(save_dir,jax.tree_map(lambda x:np.array(x),state),step=global_step,keep=3)
+
+# 수정
+import jax.tree_util
+checkpoints.save_checkpoint(save_dir, jax.tree_util.tree_map(lambda x: np.array(x), state), step=global_step, keep=3)
+
 print("Saved checkpoint to",save_dir)
 
 # ------------------