완전한 학습 루프

지금까지 배운 모델, Optax, 스케줄, 체크포인트를 합쳐 완전한 미니 학습기를 만들어 보자. 실제 프로덕션에 필요한 핵심 요소를 빠뜨리지 않은 형태야.

import jax
import jax.numpy as jnp
from jax import random
import optax
from flax import nnx
from flax import struct

# ============ 1. Model ============
class MLP(nnx.Module):
    def __init__(self, dims, *, rngs):
        self.layers = [
            nnx.Linear(dims[i], dims[i+1], rngs=rngs)
            for i in range(len(dims) - 1)
        ]

    def __call__(self, x):
        for layer in self.layers[:-1]:
            x = jax.nn.relu(layer(x))
        return self.layers[-1](x)

# ============ 2. TrainState ============
@struct.dataclass
class TrainState:
    params: any
    opt_state: any
    step: int
    key: any

# ============ 3. 초기화 ============
key = random.PRNGKey(42)
key, init_key, train_key = random.split(key, 3)

# model
model = MLP([784, 256, 128, 10], rngs=nnx.Rngs(init_key))
graphdef, params = nnx.split(model, nnx.Param)

# optimizer
schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0,
    peak_value=1e-3,
    warmup_steps=500,
    decay_steps=5000,
    end_value=1e-5,
)
optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adamw(learning_rate=schedule, weight_decay=1e-4),
)
opt_state = optimizer.init(params)

state = TrainState(
    params=params,
    opt_state=opt_state,
    step=0,
    key=train_key,
)

# ============ 4. loss 와 train step ============
def loss_fn(params, x, y, key):
    model = nnx.merge(graphdef, params)
    logits = model(x)
    log_probs = jax.nn.log_softmax(logits)
    loss = -jnp.mean(jnp.sum(y * log_probs, axis=-1))
    acc = jnp.mean(jnp.argmax(logits, axis=-1) == jnp.argmax(y, axis=-1))
    return loss, {"acc": acc}

@jax.jit
def train_step(state, batch):
    x, y = batch
    key, subkey = random.split(state.key)

    (loss, metrics), grads = jax.value_and_grad(loss_fn, has_aux=True)(
        state.params, x, y, subkey,
    )

    updates, new_opt_state = optimizer.update(grads, state.opt_state, state.params)
    new_params = optax.apply_updates(state.params, updates)

    new_state = state.replace(
        params=new_params,
        opt_state=new_opt_state,
        step=state.step + 1,
        key=key,
    )
    return new_state, loss, metrics

# ============ 5. 평가 step ============
@jax.jit
def eval_step(state, batch):
    x, y = batch
    model = nnx.merge(graphdef, state.params)
    logits = model(x)
    acc = jnp.mean(jnp.argmax(logits, axis=-1) == jnp.argmax(y, axis=-1))
    return acc

# ============ 6. 학습 루프 ============
import time

def make_batch(key, batch_size=128):
    '''더미 batch — 실제론 dataset'''
    k1, k2 = random.split(key)
    x = random.normal(k1, (batch_size, 784))
    y = jax.nn.one_hot(random.randint(k2, (batch_size,), 0, 10), 10)
    return x, y

t = time.time()
for epoch in range(10):
    # 학습
    for batch_idx in range(50):
        batch = make_batch(random.fold_in(key, epoch * 1000 + batch_idx))
        state, loss, metrics = train_step(state, batch)

    # epoch 평가
    val_batch = make_batch(random.fold_in(key, 999_999), batch_size=512)
    val_acc = eval_step(state, val_batch)

    current_lr = schedule(state.step)
    print(f"epoch {epoch:2d}  step {state.step:4d}  "
          f"loss={loss:.4f}  acc={metrics['acc']:.3f}  "
          f"val_acc={val_acc:.3f}  lr={current_lr:.6f}")

print(f"\n총 학습 시간: {time.time()-t:.1f}s")

이게 프로덕션에 적용할 수 있는 학습기의 골격이야. 추가할 만한 것:

• 실제 데이터 로더 (PyTorch DataLoader 호환 또는 grain / tf. 데이터)
• wandb / tensorboard 로깅
• 체크포인트 저장 (Track 11-5)
• multi-GPU 분산 (Track 7)
• mixed precision (Optax가 helper 제공)
• 그래디언트 accumulation (effective batch size 키우기)

이 패턴, Llama-3 70B 학습 코드도, 작은 MLP도, 기본 모양은 같아. 차이는 모델 크기와 데이터 양, 멀티호스트 분산 코드에 있어.