Update Model_torch.py

Model_torch.py

@@ -160,34 +160,43 @@ class ReLM(nn.Module):
         logits = x @ self.token_embedding.weight.T
         return logits
 
-# ===============================
-# 학습
-# ===============================
+# 모델, 옵티마이저, 스케줄러, 손실 함수
 model = ReLM(vocab_size, max_len, 128, 2).to(device)
 optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
 scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)
 loss_fn = nn.CrossEntropyLoss(ignore_index=pad_id)
 
+# 정적 그래프 컴파일
+model = torch.compile(model, mode="default")
+
+scaler = torch.cuda.amp.GradScaler()
 epochs = 1
 for epoch in range(epochs):
     model.train()
     total_loss = 0
-    for step,(x,y) in enumerate(dataloader):
-        x,y = x.to(device), y.to(device)
+    for step, (x, y) in enumerate(dataloader):
+        x, y = x.to(device), y.to(device)
         optimizer.zero_grad()
-        logits = model(x)
-        loss = loss_fn(logits.view(-1,vocab_size), y.view(-1))
-        loss.backward()
-        torch.nn.utils.clip_grad_norm_(model.parameters(),1.0)
-        optimizer.step()
+        
+        with torch.cuda.amp.autocast():  # mixed precision
+            logits = model(x)
+            loss = loss_fn(logits.view(-1, vocab_size), y.view(-1))
+        
+        scaler.scale(loss).backward()
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        scaler.step(optimizer)
+        scaler.update()
+        
         total_loss += loss.item()
         if step % 100 == 0:
             print(f"Epoch {epoch+1}, Step {step}, Loss: {loss.item():.4f}")
+    
     scheduler.step()
     print(f"Epoch {epoch+1} 완료, 평균 Loss: {total_loss/len(dataloader):.4f}")
 
 torch.save(model.state_dict(), "relm_model.pth")
-print("모델 저장 완료!")
+print("✅ 모델 저장 완료!")
 
 # ===============================
 # Top-p 샘플링 생성