GPU Training နှင့် Mixed Precision

နားလည်ထားရမယ့် အချက်

Neural network training ဟာ forward နဲ့ backward pass ထဲက matrix multiplication တွေက လွှမ်းမိုးနေပါတယ်။ GPU တွေမှာတော့ ဒီလို massively parallel arithmetic အတွက် အထူးဖန်တီးထားတဲ့ core သေးသေးလေး ထောင်နဲ့ချီ ပါဝင်နေလို့ computation ကို GPU ပေါ်ရွှေ့လိုက်ရင် real workload တွေမှာ training time ကို ဆယ်ဆလောက် လျှော့ချနိုင်ပါတယ်။ ရွှေ့တဲ့ mechanism က .to(device) ပါ — ဒါဟာ target device ပေါ်မှာ ရှိနေတဲ့ tensor (သို့) model ရဲ့ copy တစ်ခုကို ပြန်ပေးပါတယ်၊ ပြီးတော့ operation တစ်ခုမှာ ပါဝင်တဲ့ tensor တိုင်းဟာ device တူတူပေါ်မှာ ရှိနေရပါတယ် — GPU ပေါ်က model နဲ့ CPU ပေါ်က data ကို ရောသုံးလိုက်ရင် runtime device-mismatch error တက်ပါတယ်၊ ဘာဖြစ်လို့လဲဆိုတော့ PyTorch က device တွေကြား data ကို သင့်အတွက် တိတ်တဆိတ် ကူးပေးမှာ မဟုတ်လို့ပါ။ GPU ကို သုံးရုံသက်သက်ထက် mixed precision training ကတော့ ပိုပြီး တွန်းအားပေးပါတယ် — operation အများစုကို float16 (memory တစ်ဝက်၊ tensor core ပေါ်မှာ ပိုမြန်) နဲ့ run ပြီး gradient စုစည်းတာ (သို့) reduction အချို့လို numerically fragile step တွေကိုတော့ torch.autocast က float32 ထဲမှာ အလိုအလျောက် ထားပေးပါတယ်၊ GradScaler ကလည်း backward မလုပ်ခင် loss ကို scale တင်ပြီး optimizer step မလုပ်ခင် ပြန်ချပေးကာ float16 gradient သေးသေးလေးတွေ zero ဖြစ်သွားခြင်းကို ကာကွယ်ပေးပါတယ်။ ဒီ API တွေက availability ကို ကိုယ်တိုင် စစ်ထားလို့ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') နဲ့ ရေးထားတဲ့ program ဟာ CPU-only machine ပေါ်မှာ မြန်မြန်တော့ မဟုတ်ပေမယ့် မှန်ကန်စွာ run နိုင်ပါတယ်။

လက်တွေ့ scenario နဲ့ ချိတ်ကြည့်မယ်

Tutorial Platform ရဲ့ search-ranking model (သို့) lesson-recommendation model ကို click နဲ့ completion သန်းချီရှိတဲ့ interaction history အပြည့်အစုံနဲ့ train ဖို့ CPU တစ်ခုတည်းနဲ့ဆိုရင် လက်တွေ့မကျနိုင်လောက်အောင် အချိန်ကြာပါလိမ့်မယ်။ training job အတွက် GPU instance ငှားပြီး device = torch.device(...) ကို နေရာတိုင်း သုံးထားကာ .cpu() ကို hardcode မလုပ်ဘဲ device-agnostic ဖြစ်အောင် ရေးထားရင် training script တစ်ခုတည်းက laptop prototype ကနေ GPU box အထိ code ပြောင်းစရာမလိုဘဲ scale တက်နိုင်ပါတယ်။ Mixed precision ကတော့ အကြီးဆုံး model ဖြစ်တဲ့ recommendation transformer အတွက် အရေးအကြီးဆုံးဖြစ်ပြီး training time နဲ့ memory ကို ထက်ဝက်လောက် လျှော့ချပေးနိုင်ကာ team အနေနဲ့ budget တူတူထဲမှာ architecture ပိုစမ်းကြည့်နိုင်စေပါတယ် — lightweight sentiment classifier ကတော့ နည်းလမ်းနှစ်ခုစလုံးနဲ့ ပြေပြေလည်လည် train ဖြစ်ပါတယ်။

• torch.device('cuda' if torch.cuda.is_available() else 'cpu') နဲ့ device ကို detect လုပ်ပြီး model နဲ့ input tensor တိုင်းအပေါ် .to(device) ကို ခေါ်ပါ — code ထဲမှာ 'cuda' ကို hardcode လုံးဝ မလုပ်သင့်ပါဘူး။
• Operation တစ်ခုမှာ ပါဝင်တဲ့ tensor အားလုံးဟာ device တူတူပေါ်မှာ ရှိရပါမယ် — GPU နဲ့ CPU tensor ရောသုံးလိုက်ရင် အလိုအလျောက် transfer ပေးမှာ မဟုတ်ဘဲ runtime error တက်ပါတယ်။
• Mixed precision (torch.autocast + GradScaler) ဟာ operation အများစုကို float16 နဲ့ run ပြီး numerically sensitive step တွေကို float32 ထဲ အလိုအလျောက် ထားခြင်းအားဖြင့် GPU training ကို မြန်ဆန်စေပါတယ်။

အတူတူ စမ်းရေးကြည့်မယ်

import torch
import torch.nn as nn

torch.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

class TinyClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 16)
        self.fc2 = nn.Linear(16, 2)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = TinyClassifier().to(device)
sample_input = torch.randn(4, 10).to(device)

with torch.no_grad():
    output = model(sample_input)

print("Output device:", output.device)
print("Output shape:", output.shape)

# Mixed precision training sketch -- only actually uses float16 on CUDA;
# falls back safely to normal float32 compute on CPU
scaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
target = torch.randint(0, 2, (4,)).to(device)

optimizer.zero_grad()
with torch.autocast(device_type=device.type, enabled=torch.cuda.is_available()):
    logits = model(sample_input)
    loss = nn.functional.cross_entropy(logits, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

print("Training step completed. Loss:", loss.item())

Using device: cpu
Output device: cpu
Output shape: torch.Size([4, 2])
Training step completed. Loss: 0.7749
(CUDA GPU ပါတဲ့ machine ပေါ်မှာဆိုရင် ပထမ line နှစ်ကြောင်းက 'cuda' လို့ print ထွက်မှာဖြစ်ပြီး autocast ကလည်း computation တစ်စိတ်တစ်ပိုင်းကို float16 နဲ့ တကယ် run ပေးပါလိမ့်မယ်)

PyTorch Docs — Automatic Mixed Precision — Deep Learning