138 lines
4.6 KiB
Python
138 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
|
|
from airfrans_frontier.training.config import TrainingConfig
|
|
|
|
|
|
def estimate_forward_flops_per_item(model: nn.Module) -> int:
|
|
total = 0
|
|
for module in model.modules():
|
|
if isinstance(module, nn.Linear):
|
|
total += 2 * module.in_features * module.out_features
|
|
if module.bias is not None:
|
|
total += module.out_features
|
|
return int(total)
|
|
|
|
|
|
def estimate_training_compute(*, steps: int, batch_size: int, forward_flops_per_item: int) -> int:
|
|
return int(steps * batch_size * forward_flops_per_item * 3)
|
|
|
|
|
|
def checkpoint_size_bytes(run_dir: Path, name: str) -> int | None:
|
|
path = run_dir / name
|
|
if not path.is_file():
|
|
return None
|
|
return int(path.stat().st_size)
|
|
|
|
|
|
def gpu_memory_metrics(device: torch.device) -> dict[str, int | None]:
|
|
if device.type != "cuda":
|
|
return {
|
|
"gpu_memory_allocated_mb": None,
|
|
"gpu_memory_reserved_mb": None,
|
|
"gpu_memory_peak_allocated_mb": None,
|
|
}
|
|
index = device.index if device.index is not None else torch.cuda.current_device()
|
|
return {
|
|
"gpu_memory_allocated_mb": int(torch.cuda.memory_allocated(index) // (1024 * 1024)),
|
|
"gpu_memory_reserved_mb": int(torch.cuda.memory_reserved(index) // (1024 * 1024)),
|
|
"gpu_memory_peak_allocated_mb": int(torch.cuda.max_memory_allocated(index) // (1024 * 1024)),
|
|
}
|
|
|
|
|
|
def measure_training_step(
|
|
model: nn.Module,
|
|
optimizer: torch.optim.Optimizer,
|
|
features: np.ndarray,
|
|
targets: np.ndarray,
|
|
*,
|
|
batch_size: int,
|
|
steps: int,
|
|
device: torch.device,
|
|
) -> dict[str, Any]:
|
|
if steps <= 0:
|
|
raise ValueError("Calibration steps must be positive")
|
|
rng = np.random.default_rng(1729)
|
|
model.train()
|
|
if device.type == "cuda":
|
|
torch.cuda.reset_peak_memory_stats(device)
|
|
torch.cuda.synchronize(device)
|
|
started = time.perf_counter()
|
|
last_loss = 0.0
|
|
for _ in range(steps):
|
|
indices = rng.integers(0, features.shape[0], size=batch_size)
|
|
batch_features = torch.from_numpy(np.ascontiguousarray(features[indices], dtype=np.float32)).to(device)
|
|
batch_targets = torch.from_numpy(np.ascontiguousarray(targets[indices], dtype=np.float32)).to(device)
|
|
optimizer.zero_grad(set_to_none=True)
|
|
predictions = model(batch_features)
|
|
loss = F.mse_loss(predictions, batch_targets)
|
|
loss.backward()
|
|
optimizer.step()
|
|
last_loss = float(loss.detach().cpu().item())
|
|
if device.type == "cuda":
|
|
torch.cuda.synchronize(device)
|
|
elapsed = time.perf_counter() - started
|
|
return {
|
|
"calibration_steps": steps,
|
|
"step_time_seconds": elapsed / steps,
|
|
"points_per_sec": steps * batch_size / max(elapsed, 1e-12),
|
|
"last_calibration_loss": last_loss,
|
|
**gpu_memory_metrics(device),
|
|
}
|
|
|
|
|
|
def measure_validation_runtime(
|
|
model: nn.Module,
|
|
features: np.ndarray,
|
|
*,
|
|
batch_size: int,
|
|
device: torch.device,
|
|
) -> float:
|
|
model.eval()
|
|
if device.type == "cuda":
|
|
torch.cuda.synchronize(device)
|
|
started = time.perf_counter()
|
|
with torch.no_grad():
|
|
for start in range(0, features.shape[0], batch_size):
|
|
stop = min(start + batch_size, features.shape[0])
|
|
batch_features = torch.from_numpy(np.ascontiguousarray(features[start:stop], dtype=np.float32)).to(device)
|
|
model(batch_features)
|
|
if device.type == "cuda":
|
|
torch.cuda.synchronize(device)
|
|
return time.perf_counter() - started
|
|
|
|
|
|
def measure_checkpoint_size(payload: dict[str, Any]) -> int:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
path = Path(tmp) / "checkpoint.pt"
|
|
torch.save(payload, path)
|
|
return int(path.stat().st_size)
|
|
|
|
|
|
def write_calibration_report(path: str | Path, data: dict[str, Any]) -> Path:
|
|
report_path = Path(path)
|
|
report_path.parent.mkdir(parents=True, exist_ok=True)
|
|
report_path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n")
|
|
return report_path
|
|
|
|
|
|
def static_calibration_fields(config: TrainingConfig, model: nn.Module) -> dict[str, Any]:
|
|
forward_flops = estimate_forward_flops_per_item(model)
|
|
return {
|
|
"estimated_forward_flops_per_item": forward_flops,
|
|
"estimated_train_flops": estimate_training_compute(
|
|
steps=config.optim.steps,
|
|
batch_size=config.data.batch_size,
|
|
forward_flops_per_item=forward_flops,
|
|
),
|
|
}
|