File size: 4,121 Bytes
7f316fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | import math
from contextlib import nullcontext
from typing import Optional
import torch
from flow_matching.loss import MixturePathGeneralizedKL, EditFlowsLoss
from flow_matching.path import ProbPath
from omegaconf.dictconfig import DictConfig
from torch import nn, Tensor
from torch.cuda.amp import GradScaler
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader
from utils.logging import TrainLogger
from .flow import SourceDistribution
from .state import TrainState
from ..model.utils import build_z0_z1_with_alignment, remove_eps
from dataclasses import dataclass
from typing import List, Tuple, Optional
import torch
from torch import Tensor
import pdb
def step(
state: TrainState,
loss_fn: nn.Module, # EditFlowsLoss
path: ProbPath, # EditFlowsPathAdapter (exposes .scheduler)
scaler: GradScaler,
iterator: DataLoader,
device: torch.device,
source_distribution: SourceDistribution,
logger: TrainLogger,
training: bool,
optim_params: Optional[DictConfig],
pad_id: int,
bos_id: int,
eos_id: int,
) -> Tensor:
assert (training and (optim_params is not None)) or (not training)
state.train() if training else state.eval()
batch = next(iterator)
# x_1 = pad_sequence(batch['input_ids'], batch_first=True, padding_value=pad_id).to(device)
x_1 = torch.tensor(batch["input_ids"]).to(device)
B = x_1.shape[0]
# === Source & time ===
with torch.no_grad():
eps_id = getattr(path, "eps_id", -1)
allowed_tokens = torch.tensor([tok for tok in source_distribution._allowed_tokens if tok != eps_id]).to(device)
x_0 = source_distribution.sample_x0_from_x1(x_1, pad_id=pad_id, allowed_tokens=allowed_tokens, scale_size=2, bos_id = bos_id, eos_id = eos_id)
t = torch.rand(B, device=device)
sched = path.scheduler(t)
precomputed_weight = sched.d_alpha_t / sched.sigma_t # (B,)
z_0, z_1 = build_z0_z1_with_alignment(x_0, x_1, eps_id, pad_id, bos_id, eos_id, p_optimal=0.6)
z_t = path.sample(z_0, z_1, t=t)
x_t, mask = remove_eps(z_t, eps_id, pad_id)
ctx = torch.amp.autocast('cuda', dtype=torch.float16) if training else torch.no_grad()
with ctx:
# pdb.set_trace()
lam_ins, logits_ins, lam_del, lam_sub, logits_sub = state.model(x_t=x_t, mask=mask,t=t)
loss = loss_fn(lam_ins, logits_ins, lam_del, lam_sub, logits_sub,
z_t, z_1, x_t, mask, precomputed_weight, eps_id, bos_id, eos_id)
if training:
optimization_step(
state=state,
loss=loss,
scaler=scaler,
optim_params=optim_params,
logger=logger,
)
return loss.detach()
def _get_lr(lr: float, step: int, warmup: int, n_iters: int, eta_min_ratio: float):
if step < warmup:
# Linear warmup
return lr * (step / warmup)
else:
# Cosine annealing
total_steps = n_iters
eta_min = eta_min_ratio * lr
cosine_decay = 0.5 * (
1 + math.cos(math.pi * (step - warmup) / (total_steps - warmup))
)
return eta_min + (lr - eta_min) * cosine_decay
def optimization_step(
state: TrainState,
scaler: GradScaler,
loss: Tensor,
optim_params: DictConfig,
logger: TrainLogger,
) -> None:
scaler.scale(loss).backward()
scaler.unscale_(state.optimizer)
lr = _get_lr(
lr=optim_params.lr,
step=state.step,
warmup=optim_params.warmup,
n_iters=optim_params.n_iters,
eta_min_ratio=optim_params.eta_min_ratio,
)
# Update learning rate in optimizer
for g in state.optimizer.param_groups:
g["lr"] = lr
if state.step % optim_params.log_lr_every == 0:
logger.log_lr(value=lr, step=state.step)
if optim_params.grad_clip >= 0:
torch.nn.utils.clip_grad_norm_(
state.model.parameters(), max_norm=optim_params.grad_clip
)
scaler.step(state.optimizer)
scaler.update()
state.optimizer.zero_grad()
|