Update Mo.py

Mo.py

@@ -6,7 +6,11 @@ import requests
 from tensorflow import keras
 from tensorflow.keras import layers
 import tensorflow.keras.backend as K
-
+# ===============================
+from tensorflow.keras import mixed_precision
+policy = mixed_precision.Policy('mixed_float16')  # fp16
+mixed_precision.set_global_policy(policy)
+print("✅ Mixed precision 적용:", policy)
 print('1')
 tf.get_logger().setLevel("ERROR")
 SEED = 42
@@ -63,8 +67,8 @@ unk_id = sp.piece_to_id("<unk>")
 vocab_size = sp.get_piece_size()
 print(f"✅ Vocabulary size: {vocab_size}")
 
-max_len = 256
-batch_size = 128
+max_len = 200
+batch_size = 96
 
 def text_to_ids(text):
     return sp.encode(text, out_type=int)
@@ -169,20 +173,22 @@ class MHLA(layers.Layer):
 class Lo(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
-        self.d = layers.Dense(64, activation='silu')
-        self.w = layers.Dense(d_model)
-        self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
+        self.d = layers.Dense(64, activation='silu', dtype='float16')  # fp16 연산
+        self.w = layers.Dense(d_model, dtype='float16')               # fp16 연산
+        self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')  # fp32
 
     def call(self, x):
         p = self.d(x)
         p = self.w(p)
-        return self.norm(p) + x
+        p = self.norm(p)  # fp32
+        return tf.cast(p, x.dtype) + x  # 다시 fp16로 맞춰서 Add
+
 
 class Block(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
         self.lou = MHLA(d_model, 8)
-        self.glu = SwiGLU(d_model, 1154)
+        self.glu = SwiGLU(d_model, 1048)
         self.lo = Lo(d_model)
 
     def call(self, x):
@@ -193,10 +199,10 @@ class Block(layers.Layer):
 class LaSLM(tf.keras.Model):
     def __init__(self, vocab_size, max_seq_len, d_model, n_layers, dropout_rate=0.1):
         super().__init__()
-        self.token_embedding = layers.Embedding(vocab_size, d_model)
-        self.pos_embedding = layers.Embedding(max_seq_len, d_model)
+        self.token_embedding = layers.Embedding(vocab_size, d_model, dtype=policy.compute_dtype)
+        self.pos_embedding = layers.Embedding(max_seq_len, d_model, dtype=policy.compute_dtype)
         self.blocks = [Block(d_model) for _ in range(n_layers)]
-        self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype="float32")
+        self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype='float32')  # ln_f는 fp32
 
     def call(self, x, training=False):
         batch_size, seq_len = tf.shape(x)[0], tf.shape(x)[1]
@@ -207,7 +213,7 @@ class LaSLM(tf.keras.Model):
         x = self.ln_f(x)
         embedding_matrix = tf.cast(self.token_embedding.embeddings, x.dtype)
         logits = tf.matmul(x, embedding_matrix, transpose_b=True)
-        return tf.cast(logits, tf.float32)
+        return tf.cast(logits, tf.float32)  # loss 계산을 위해 fp32로 변환
 
 def smoothed_loss_keras(y_true, y_pred, eps=0.1):
     y_true = tf.cast(y_true, tf.int32)