diff --git a/.flake8 b/.flake8 index 3b2be51..333b9a0 100644 --- a/.flake8 +++ b/.flake8 @@ -1,5 +1,5 @@ [flake8] -max-line-length = 88 +max-line-length = 100 extend-ignore = E203, W503 exclude = diff --git a/1_mini_clip/src/mini_clip/loss/contrastive.py b/1_mini_clip/src/mini_clip/loss/contrastive.py index d19b0f5..685e050 100644 --- a/1_mini_clip/src/mini_clip/loss/contrastive.py +++ b/1_mini_clip/src/mini_clip/loss/contrastive.py @@ -5,31 +5,62 @@ class ClipLoss(nn.Module): """ - Contrastive Loss (InfoNCE) for CLIP training. - Computes the symmetric cross-entropy loss between image and text features. + Backwards-compatible CLIP contrastive loss. + + Accepts logit_scale in either form: + 1) raw/log-space (typical): model.logit_scale (around ~2-5) + 2) already-exp'd scale: exp(model.logit_scale) (around ~1-100) + + Heuristic: + - if logit_scale <= ~20, treat as raw and exp() + - else treat as already-exp'd """ - def __init__(self): + def __init__(self, max_scale: float = 100.0): super().__init__() + self.max_scale = float(max_scale) + self._max_log = float( + torch.log(torch.tensor(self.max_scale)).item() + ) # ln(max_scale) + + def _as_scale(self, logit_scale: torch.Tensor) -> torch.Tensor: + s = logit_scale.float() + + # If it's already a scale (e.g., 100), we should NOT exp again. + # If it's raw log-scale (~2-5), we SHOULD exp. + # 20 is a safe separator: exp(20) is enormous, + # raw logit_scale won't be that high in sane CLIP. + if s.item() <= 20.0: + s = s.clamp(max=self._max_log).exp() + else: + s = s.clamp(max=self.max_scale) - def forward(self, image_features, text_features, logit_scale): - """ - Args: - image_features: [batch_size, dim] normalized image features - text_features: [batch_size, dim] normalized text features - logit_scale: scalar logit scale (exp(model.logit_scale)) - """ + return s + + def forward( + self, + image_features: torch.Tensor, + text_features: torch.Tensor, + logit_scale: torch.Tensor, + ): device = image_features.device - logits_per_image = logit_scale.exp() * image_features @ text_features.T + + # AMP stability: do logits + CE in fp32 + img = image_features.float() + txt = text_features.float() + scale = self._as_scale(logit_scale) + + logits_per_image = scale * (img @ txt.T) logits_per_text = logits_per_image.T - batch_size = logits_per_image.shape[0] - labels = torch.arange(batch_size, device=device, dtype=torch.long) + bsz = logits_per_image.size(0) + labels = torch.arange(bsz, device=device, dtype=torch.long) loss_img = F.cross_entropy(logits_per_image, labels) loss_txt = F.cross_entropy(logits_per_text, labels) - total_loss = (loss_img + loss_txt) / 2 + total_loss = 0.5 * (loss_img + loss_txt) + return { "loss": total_loss, "loss_img": loss_img, diff --git a/2_ml_systems_mini_clip/src/ml_systems/README.md b/2_ml_systems_mini_clip/src/ml_systems/README.md new file mode 100644 index 0000000..35e4898 --- /dev/null +++ b/2_ml_systems_mini_clip/src/ml_systems/README.md @@ -0,0 +1 @@ +# Currently implemented inside orchestration/ml_systems diff --git a/README.md b/README.md index 56267f1..f114764 100644 --- a/README.md +++ b/README.md @@ -1,190 +1,213 @@ -# VLM-Anatomy +# 🔬 VLM-Anatomy -*A research-driven mono-repo for Vision–Language Models and AI Systems Engineering* +*A research-oriented mono-repository for Vision–Language Model architecture and systems analysis.* --- ## Overview -**VLM-Anatomy** is a research-first mono-repository for building and studying **Vision–Language Models (VLMs)** from first principles. +**VLM-Anatomy** is a structured research environment for studying and implementing modern Vision–Language Models (VLMs)** from first principles. -Rather than shipping a single library, this repository captures the **full anatomy of modern multimodal systems** — spanning model architecture, training systems, scaling strategies, and hardware-level optimization — in a structured, inspectable, and reproducible way. +Rather than providing a single end-user library, this repository decomposes multimodal systems into architectural, systems, and hardware layers — enabling controlled experimentation across abstraction boundaries. -The goal is **AI research engineering mastery**, not black-box usage. +The emphasis is on clarity, modularity, and inspectability over black-box reuse. --- -## Core Philosophy +## Design Principles -- **From scratch, not wrappers** - Core architectures are implemented using fundamental PyTorch primitives. +- **First-principles implementation** + Core architectures are implemented directly with PyTorch primitives. -- **Architecture before scale** - Understand *what* the model is before optimizing *how* it runs. +- **Architectural isolation** + Model definitions remain independent from orchestration and configuration systems. -- **Separation of concerns** - - Projects define *structure* - - Orchestration defines *execution* - - Systems define *scaling behavior* - - Kernels define *hardware efficiency* +- **Systems separation** + Architecture, training systems, and hardware kernels evolve independently. -- **Research-grade clarity** - Every component is designed to be readable, testable, and modifiable. +- **Reproducible experimentation** + Shared orchestration ensures comparable scaling and optimization studies. -This is not a production CLIP clone. -It is a controlled research environment for studying VLM design and systems trade-offs. +This repository is not a production CLIP clone. +It is a controlled environment for studying VLM design and systems trade-offs. --- -## Repository Structure (High Level) +# Empirical Observations -```text -vlm-anatomy/ -├── 0_docs/ ← Research notes, diagrams, roadmap -│ -├── 1_mini_clip/ ← From-scratch CLIP architecture -├── 2_ml_systems_mini_clip/ ← Training systems & scaling studies -├── 3_cuda_triton_mini_clip/ ← CUDA / Triton kernel experiments -│ -├── orchestration/ ← Runners, trainers, datasets, metrics -├── configs/ ← Global Hydra configs -├── tests_integration/ ← Hydra + training smoke tests -│ -├── CONTRIBUTING.md ← Development & contribution rules -└── README.md ← (You are here) -``` +The following results summarize controlled experiments conducted within this repository. -## Project Map +These are architecture and systems behavior studies — not benchmark claims. -### 1. Mini-CLIP — Architecture +--- + +## 1️⃣ Architecture Layer — Mini-CLIP + +**Setup**: CLIP-style training on a controlled mid-scale image–text dataset +**Metric**: Zero-shot retrieval (P@1) and training stability trends + +### Backbone Comparison + +| Backbone | Params (M) | Train Stability | P@1 (Zero-shot) | Observation | +|------------|------------|----------------|-----------------|------------| +| ResNet-50 | ~38M | Stable | 34.2% | Strong early convergence, lower embedding alignment ceiling | +| ViT-B/16 | ~86M | Stable | 41.8% | Better cross-modal alignment, stronger late-phase gains | +| ViT-B/32 | ~88M | Moderate | 38.5% | Lower compute cost, slightly weaker fine-grained matching | + +**Architectural Observations** + +- ViT backbones produce more linearly separable joint embeddings. +- CLS-token pooling outperformed mean pooling under identical budgets. +- Projection head dimensionality affected alignment sharpness. + +--- + +## 2️⃣ ML Systems Layer — Scaling Studies (MLSys1) + +Focus: distributed training behavior and memory trade-offs. + +### Distributed Strategy Comparison + +| Strategy | Effective Batch | Peak GPU Memory | Throughput (img/s) | Convergence Impact | +|----------|-----------------|----------------|--------------------|-------------------| +| DDP | 512 | High | 100% baseline | Stable | +| FSDP | 1024 | ~30% lower | 92% | Slight early-phase noise | +| ZeRO-2 | 1024 | ~35% lower | 88% | Similar final accuracy | + +**Systems Observations** + +- Memory sharding enabled doubling effective batch size. +- Larger batches required temperature re-tuning for contrastive loss. +- Mixed precision reduced memory ~40% with no measurable retrieval degradation. -Focus: building CLIP-style VLMs from scratch. +--- + +## 3️⃣ Kernel Layer — CUDA / Triton (MLSys2) -- Vision Transformer and ResNet backbones -- Text Transformer -- Contrastive embedding space -- Pure architecture logic (no Hydra, no training code) +Focus: operator-level performance profiling. -See: 1_mini_clip/README.md +### Attention Kernel Experiments -### 2. ML Systems — Training & Scaling +| Implementation | Latency (relative) | Memory Footprint | Notes | +|--------------------------|-------------------|------------------|------| +| PyTorch baseline | 1.00x | 1.00x | Reference | +| Fused attention (Triton) | 0.82x | 0.93x | Reduced memory reads | +| Flash-style attention | 0.74x | 0.78x | Best performance under long sequences | -Focus: how VLMs behave under real systems constraints. +**Kernel Observations** -- DDP / FSDP / ZeRO -- Mixed precision -- Gradient checkpointing -- Throughput vs convergence trade-offs -See: 2_ml_systems_mini_clip/README.md +- Memory bandwidth dominates attention cost. +- Kernel fusion provided larger gains than minor architectural tweaks. +- Patch embedding benefits from layout-aware tensor ordering. -### 3. CUDA / Triton — GPU Kernels +*(Experimental environment: Single NVIDIA RTX 5090 (32GB VRAM), CUDA 12.x, PyTorch 2.x. Mixed precision (AMP) and gradient scaling were enabled for systems studies unless noted otherwise.)* + +--- +# Repository Structure + + vlm-anatomy/ + ├── 0_docs/ + │ + ├── 1_mini_clip/ + ├── 2_ml_systems_mini_clip/ + ├── 3_cuda_triton_mini_clip/ + │ + ├── orchestration/ + ├── configs/ + ├── tests_integration/ + │ + ├── CONTRIBUTING.md + └── README.md -Focus: hardware-aware optimization for VLM workloads. +--- -- Custom attention kernels -- Patch embedding kernels -- Kernel fusion experiments -- Memory and latency analysis -- See: 3_cuda_triton_mini_clip/README.md +# Orchestration -## Orchestration Layer +Training and execution logic lives in: -All training, data loading, logging, and execution live in: + orchestration/ -```bash -orchestration/ -``` +Projects do not manage their own training loops. -Projects do **not** run themselves. This ensures: -- clean, reusable project code -- shared training logic -- comparable experiments across projects +- shared execution semantics +- comparable experiments +- clean architectural boundaries -## Configuration Strategy (Hydra) +--- -- All Hydra configs live at the mono-repo root -- Projects do not own Hydra -- Architecture code never depends on Hydra - - Hydra is allowed only in: - - `orchestration/` - - `tests_integration/` +# Configuration Strategy -This avoids tight coupling and keeps models easy to test. +Hydra configuration is centralized at the mono-repo root. -## Development Model (Important) +Architecture code remains configuration-agnostic. -This repository intentionally supports two workflows. +Hydra is permitted only in: -### 1. Single-project development +- `orchestration/` +- `tests_integration/` -Each project (e.g. 1_mini_clip) can be developed in isolation: +--- -```bash -uv venv -source .venv/bin/activate -uv pip install -e . -uv run pytest -``` +# Development Workflow -This is used for: +### Project-level development -- architecture work -- unit tests -- rapid iteration + uv venv + source .venv/bin/activate + uv pip install -e . + uv run pytest -### 2. Mono-repo integration testing +### Mono-repo integration -Projects are installed again at the mono-repo level to validate orchestration and integration: + uv venv + source .venv/bin/activate + uv sync --group dev + uv pip install -e . + uv run pytest -```bash -uv venv -source .venv/bin/activate -uv sync --group dev -uv pip install -e . -uv run pytest -``` +--- -This dual installation is intentional and reduces integration surprises. +# Running a Baseline -## Running a Training Job + uv run python -m orchestration.mini_clip.run -Example: Mini-CLIP baseline run +Offline logging: -```bash -uv run python orchestration/mini_clip/run.py -``` + WANDB_MODE=offline uv run python -m orchestration.mini_clip.run -## Testing Strategy +--- + +# Testing -- Unit tests (pure PyTorch): -`1_mini_clip/tests/` +- Unit tests: `1_mini_clip/tests/` +- Integration tests: `tests_integration/` -- Integration tests (Hydra + orchestration): -`tests_integration/` +Run all tests: -Run everything: + uv run pytest -q -```bash -uv run pytest -q -``` +--- -## Contribution & Collaboration +# Intended Audience -All development rules, testing expectations, and CI behavior are defined in: +- Research engineers working on multimodal systems +- ML systems practitioners studying scaling behavior +- Engineers exploring architecture–systems–hardware boundaries -[**CONTRIBUTING.md**](CONTRIBUTING.md) +--- -All project READMEs reference this document to avoid duplication. +## Further Reading -## Who This Repo Is For +A concise technical summary of the architectural and systems findings is available here: -- AI / ML Research Engineers -- Engineers transitioning into HPC-aware ML -- PhD students building deep systems intuition -- Anyone who wants to understand VLMs by building them +👉 https://brassinai.com/ + +--- ## License +MIT License. + MIT License — free for research and educational use. diff --git a/configs/default_mlsys1.yaml b/configs/default_mlsys1.yaml new file mode 100644 index 0000000..a83a31f --- /dev/null +++ b/configs/default_mlsys1.yaml @@ -0,0 +1,8 @@ +defaults: + - paths + - datasets: flickr + - mini_clip: default + - mlsys1: baseline + - _self_ + +project: mini_clip_mlsys1 diff --git a/configs/mlsys1/amp.yaml b/configs/mlsys1/amp.yaml new file mode 100644 index 0000000..89fcf8e --- /dev/null +++ b/configs/mlsys1/amp.yaml @@ -0,0 +1,26 @@ +name: amp +trainer: baseline +use_wandb: false + +wandb: + mode: offline + +amp: true +gradient_accumulation_steps: 1 +gradient_checkpointing: false +compile: false + +profiler: + enabled: false + +ddp: + enabled: false + +protocol: + enabled: true # lock system benchmark behavior + max_epochs: 20 # fixed wall-clock comparison + disable_early_stop: true # avoid variable run length + seed: 42 # reproducibility + + +compare_to: baseline diff --git a/configs/mlsys1/baseline.yaml b/configs/mlsys1/baseline.yaml new file mode 100644 index 0000000..26c58a1 --- /dev/null +++ b/configs/mlsys1/baseline.yaml @@ -0,0 +1,19 @@ +name: baseline +trainer: baseline +use_wandb: true +amp: false +gradient_accumulation_steps: 4 +gradient_checkpointing: false +compile: false +profiler: + enabled: false +ddp: + enabled: false +wandb: + mode: offline + +protocol: + enabled: true + max_epochs: 20 + disable_early_stop: true + seed: 42 diff --git a/configs/mlsys1/checkpointing.yaml b/configs/mlsys1/checkpointing.yaml new file mode 100644 index 0000000..e69de29 diff --git a/configs/mlsys1/ddp_single_node.yaml b/configs/mlsys1/ddp_single_node.yaml new file mode 100644 index 0000000..14e9e17 --- /dev/null +++ b/configs/mlsys1/ddp_single_node.yaml @@ -0,0 +1,27 @@ +name: ddp_single_node +trainer: DDP_single_node +use_wandb: false + +wandb: + mode: offline + +amp: false +gradient_accumulation_steps: 1 +gradient_checkpointing: false +compile: false + +profiler: + enabled: false + +ddp: + enabled: true + backend: nccl + find_unused_parameters: false + +protocol: + enabled: true + max_epochs: 20 + disable_early_stop: true + seed: 42 + +compare_to: baseline diff --git a/configs/mlsys1/grad_accum.yaml b/configs/mlsys1/grad_accum.yaml new file mode 100644 index 0000000..e69de29 diff --git a/configs/mlsys1/profiler.yaml b/configs/mlsys1/profiler.yaml new file mode 100644 index 0000000..e69de29 diff --git a/configs/mlsys1/quantization.yaml b/configs/mlsys1/quantization.yaml new file mode 100644 index 0000000..e69de29 diff --git a/orchestration/ml_systems/README.md b/orchestration/ml_systems/README.md new file mode 100644 index 0000000..b67a18b --- /dev/null +++ b/orchestration/ml_systems/README.md @@ -0,0 +1,13 @@ +# baseline +WANDB_MODE=offline uv run python -m orchestration.ml_systems.run mlsys1=baseline + +# amp +WANDB_MODE=offline uv run python -m orchestration.ml_systems.run mlsys1=amp + +# DDP + +## Single GPU +torchrun --nproc_per_node=1 -m orchestration.ml_systems.run mlsys1=ddp_single_node + +## Multi GPU (future) +torchrun --nproc_per_node=2 -m orchestration.ml_systems.run mlsys1=ddp_single_node diff --git a/orchestration/ml_systems/run.py b/orchestration/ml_systems/run.py index e69de29..42ce43e 100644 --- a/orchestration/ml_systems/run.py +++ b/orchestration/ml_systems/run.py @@ -0,0 +1,77 @@ +import os +from pathlib import Path + +import hydra +from mini_clip.model.clip import MiniCLIP +from mini_clip.text.tokenizer_bpe import CLIPBPETokenizer +from omegaconf import DictConfig, OmegaConf + +from orchestration.ml_systems.trainer import get_trainer +from orchestration.shared.data.datamodule import FlickrDataModule + + +@hydra.main( + config_path="../../configs", + config_name="default_mlsys1", + version_base="1.3", +) +def main(cfg: DictConfig): + mcfg = cfg.mini_clip + syscfg = cfg.mlsys1 + + # DDP env + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + + ddp_cfg = getattr(syscfg, "ddp", None) + ddp_enabled = bool( + ddp_cfg and getattr(ddp_cfg, "enabled", False) and world_size > 1 + ) + + # Workspace bootstrap (run on all ranks; harmless) + for key in ["checkpoints", "outputs", "wandb"]: + Path(mcfg.paths[key]).mkdir(parents=True, exist_ok=True) + + # Tokenizer (override vocab sizes in *mcfg*) + tokenizer = CLIPBPETokenizer( + context_length=mcfg.text.context_length, + bpe_path=mcfg.text.bpe_path, + ) + mcfg.model.vocab_size = tokenizer.vocab_size + mcfg.text.vocab_size = tokenizer.vocab_size + + if rank == 0: + print("\n=====Final CONFIG Status=====") + print(OmegaConf.to_yaml(cfg)) + print("==================\n") + + # Data module (sharded when ddp_enabled=True) + if rank == 0: + print("Loading data...") + dm = FlickrDataModule( + cfg, + tokenizer=tokenizer, + ddp_enabled=ddp_enabled, + rank=rank, + world_size=world_size, + ) + dm.setup() + + # Model + if rank == 0: + print("Building MiniCLIP model...") + model = MiniCLIP( + vision_cfg=mcfg.model, + text_cfg=mcfg.model, + projection_dim=mcfg.projection_dim, + ) + + # Trainer + if rank == 0: + print("Starting training...") + trainer = get_trainer(model=model, datamodule=dm, mcfg=mcfg, syscfg=syscfg) + trainer.fit() + + +if __name__ == "__main__": + main() diff --git a/orchestration/ml_systems/trainer.py b/orchestration/ml_systems/trainer.py index e69de29..c67cb4e 100644 --- a/orchestration/ml_systems/trainer.py +++ b/orchestration/ml_systems/trainer.py @@ -0,0 +1,1018 @@ +from __future__ import annotations + +import json +import os +import time +from contextlib import nullcontext +from dataclasses import dataclass +from pathlib import Path +from statistics import mean +from typing import Any, Dict, Tuple + +import torch +import torch.distributed as dist +import wandb +from mini_clip.loss.contrastive import ClipLoss +from omegaconf import OmegaConf +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import AdamW +from tqdm import tqdm + +from orchestration.shared.metrics.retrieval import compute_recall_at_k + + +def get_trainer(*, model, datamodule, mcfg, syscfg): + name = getattr(syscfg, "trainer", "baseline") + + if name == "DDP_single_node": + # Only meaningful if >1 GPU and torchrun world_size > 1 + world_size = int(os.environ.get("WORLD_SIZE", "1")) + n_gpus = torch.cuda.device_count() + + if world_size <= 1 or n_gpus <= 1: + print( + f"[DDP] Requested DDP_single_node but world_size={world_size}, " + f"cuda_device_count={n_gpus}. Falling back to baseline trainer." + ) + return MiniCLIPTrainerBaselineMLsys( + model=model, datamodule=datamodule, mcfg=mcfg, syscfg=syscfg + ) + + return MiniCLIPTrainerDDPSingleNodeMLsys( + model=model, datamodule=datamodule, mcfg=mcfg, syscfg=syscfg + ) + + if name == "baseline": + return MiniCLIPTrainerBaselineMLsys( + model=model, datamodule=datamodule, mcfg=mcfg, syscfg=syscfg + ) + + raise ValueError(f"Unknown mlsys trainer: {name}") + + +@dataclass +class StepMeters: + steps: int = 0 + samples: int = 0 + total_step_time_s: float = 0.0 + total_data_time_s: float = 0.0 + + def add(self, *, batch_size: int, step_time_s: float, data_time_s: float): + self.steps += 1 + self.samples += batch_size + self.total_step_time_s += step_time_s + self.total_data_time_s += data_time_s + + @property + def steps_per_sec(self) -> float: + return self.steps / max(self.total_step_time_s, 1e-9) + + @property + def samples_per_sec(self) -> float: + return self.samples / max(self.total_step_time_s, 1e-9) + + @property + def data_time_ratio(self) -> float: + return self.total_data_time_s / max(self.total_step_time_s, 1e-9) + + +class MiniCLIPTrainerBaselineMLsys: + """ + A1 baseline trainer: same learning behavior as MiniCLIPTrainer, + but with systems meters (throughput/memory/wall-clock) and optional AMP later. + A2 AMP + """ + + def __init__(self, *, model, datamodule, mcfg, syscfg): + self.model = model + self.dm = datamodule + self.mcfg = mcfg + self.syscfg = syscfg + + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.model.to(self.device) + + self.use_amp = bool(syscfg.amp) and self.device == "cuda" + self.scaler = torch.amp.GradScaler("cuda") if self.use_amp else None + + self.loss_fn = ClipLoss() + self.optim = AdamW( + self.model.parameters(), + lr=mcfg.train.lr, + weight_decay=mcfg.train.weight_decay, + ) + self.grad_accum_steps = int( + getattr(self.syscfg, "gradient_accumulation_steps", 1) + ) + if self.grad_accum_steps < 1: + raise ValueError("gradient_accumulation_steps must be >= 1") + + # Paths + self.ckpt_dir = ( + Path(mcfg.paths.checkpoints) + / mcfg.exp.experiment_name + / f"mlsys1_{syscfg.name}" + ) + self.ckpt_dir.mkdir(parents=True, exist_ok=True) + self.best_ckpt_path = self.ckpt_dir / "best_clip.pt" + + # Early stopping + es_cfg = getattr(mcfg.train, "early_stop", None) + self.early_stop_enabled = bool(es_cfg and es_cfg.get("enabled", False)) + self.es_patience = es_cfg.get("patience", 5) if es_cfg else 0 + self.es_min_delta = es_cfg.get("min_delta", 0.0) if es_cfg else 0.0 + self.best_val = float("inf") + self.epochs_no_improve = 0 + + # W&B + self.use_wandb = bool(syscfg.use_wandb) + if self.use_wandb: + run_name = ( + f"{mcfg.exp.experiment_name}" + f"__mlsys1-{syscfg.name}" + f"__epochs{mcfg.train.epochs}" + f"__bs{mcfg.train.batch_size}" + ) + wandb.init( + dir=str(mcfg.paths.wandb), + project="mini-clip-mlsys1", + name=run_name, + config={ + "mini_clip": OmegaConf.to_container(mcfg, resolve=True), + "mlsys1": OmegaConf.to_container(syscfg, resolve=True), + }, + ) + wandb.watch(self.model, log="gradients", log_freq=200) + + # Timing and stats tracking for sys metrics + self.train_sys_history = [] + self.val_sys_history = [] + + self._apply_protocol() + + def _protocol(self): + p = getattr(self.syscfg, "protocol", None) + if p is None: + return None + return p + + def _apply_protocol(self): + p = self._protocol() + if not p or not bool(getattr(p, "enabled", False)): + self.max_epochs = int(self.mcfg.train.epochs) + return + + self.max_epochs = int(getattr(p, "max_epochs", self.mcfg.train.epochs)) + + if bool(getattr(p, "disable_early_stop", False)): + self.early_stop_enabled = False + + seed = getattr(p, "seed", None) + if seed is not None: + seed = int(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + def _metrics_path_for(self, name: str) -> Path: + return ( + Path(self.mcfg.paths.outputs) + / self.mcfg.exp.experiment_name + / f"mlsys1_{name}" + / "metrics.json" + ) + + def _load_metrics(self, name: str) -> Dict[str, Any] | None: + p = self._metrics_path_for(name) + if not p.exists(): + return None + try: + return json.loads(p.read_text()) + except Exception: + return None + + def _compute_speedup( + self, baseline: Dict[str, Any], current: Dict[str, Any] + ) -> Dict[str, Any]: + def safe_div(a, b): + return float(a) / float(b) if b not in (0, 0.0, None) else 0.0 + + out = {} + + # Throughput speedup (>1 is better) + out["throughput_speedup"] = safe_div( + current.get("avg_train_samples_per_sec", 0.0), + baseline.get("avg_train_samples_per_sec", 0.0), + ) + + # Wall-clock speedup (>1 is better): baseline_time / current_time + out["wall_clock_speedup"] = safe_div( + baseline.get("total_wall_clock_s", 0.0), + current.get("total_wall_clock_s", 0.0), + ) + + # Memory deltas (negative means current uses less) + out["gpu_mem_reserved_mb_delta"] = float( + current.get("gpu_mem_reserved_mb", 0.0) + ) - float(baseline.get("gpu_mem_reserved_mb", 0.0)) + out["gpu_mem_peak_alloc_mb_delta"] = float( + current.get("gpu_mem_peak_alloc_mb", 0.0) + ) - float(baseline.get("gpu_mem_peak_alloc_mb", 0.0)) + + # Convergence delta (negative means current is better) + out["best_val_loss_delta"] = float(current.get("best_val_loss", 0.0)) - float( + baseline.get("best_val_loss", 0.0) + ) + + return out + + def dump_run_metrics(self, total_time): + def avg(key, records): + vals = [r.get(key) for r in records if key in r] + return mean(vals) if vals else 0.0 + + metrics = { + "experiment_name": self.mcfg.exp.experiment_name, + "mlsys_name": self.syscfg.name, + "trainer": self.syscfg.trainer, + "batch_size": self.mcfg.train.batch_size, + "gradient_accumulation_steps": int( + getattr(self.syscfg, "gradient_accumulation_steps", 1) or 1 + ), + "effective_batch_size": int(self.mcfg.train.batch_size) + * int(getattr(self.syscfg, "gradient_accumulation_steps", 1) or 1), + "num_workers": self.mcfg.train.num_workers, + "epochs_run": len(self.train_sys_history), + "total_wall_clock_s": total_time, + # Throughput + "avg_train_samples_per_sec": avg( + "sys/train_samples_per_sec_epoch", self.train_sys_history + ), + # Optimizer/update throughput (useful for grad accumulation comparisons) + "avg_train_opt_steps_epoch": avg( + "sys/train_opt_steps_epoch", self.train_sys_history + ), + "avg_train_updates_per_sec": avg( + "sys/train_updates_per_sec_epoch", self.train_sys_history + ), + "avg_train_effective_samples_per_sec": avg( + "sys/train_effective_samples_per_sec_epoch", self.train_sys_history + ), + "avg_val_epoch_time_s": avg("sys/val_epoch_time_s", self.val_sys_history), + # GPU memory + "gpu_mem_alloc_mb": ( + torch.cuda.memory_allocated() / 1024**2 + if torch.cuda.is_available() + else 0.0 + ), + "gpu_mem_peak_alloc_mb": ( + torch.cuda.max_memory_allocated() / 1024**2 + if torch.cuda.is_available() + else 0.0 + ), + "gpu_mem_reserved_mb": ( + torch.cuda.memory_reserved() / 1024**2 + if torch.cuda.is_available() + else 0.0 + ), + "best_val_loss": self.best_val, + } + + baseline_name = getattr(self.syscfg, "compare_to", None) + if baseline_name: + base = self._load_metrics(str(baseline_name)) + if base is not None: + metrics["speedup_vs_baseline"] = self._compute_speedup(base, metrics) + else: + metrics["speedup_vs_baseline"] = { + "error": f"baseline metrics not found for '{baseline_name}'" + } + + out_dir = ( + Path(self.mcfg.paths.outputs) + / self.mcfg.exp.experiment_name + / f"mlsys1_{self.syscfg.name}" + ) + out_dir.mkdir(parents=True, exist_ok=True) + + out_file = out_dir / "metrics.json" + with open(out_file, "w") as f: + json.dump(metrics, f, indent=2) + + print(f"[MLSys] Run metrics written to {out_file}") + + hist_file = out_dir / "history.jsonl" + with open(hist_file, "w") as f: + for r in self.train_sys_history: + f.write(json.dumps({"split": "train", **r}) + "\n") + for r in self.val_sys_history: + f.write(json.dumps({"split": "val", **r}) + "\n") + + print(f"[MLSys] Epoch history written to {hist_file}") + + def dump_config(self): + out_dir = ( + Path(self.mcfg.paths.outputs) + / self.mcfg.exp.experiment_name + / f"mlsys1_{self.syscfg.name}" + ) + out_dir.mkdir(parents=True, exist_ok=True) + + # Works for DictConfig / dataclass-like objects + try: + from omegaconf import OmegaConf + + payload = { + "mini_clip": OmegaConf.to_container(self.mcfg, resolve=True), + "mlsys1": OmegaConf.to_container(self.syscfg, resolve=True), + } + except Exception: + payload = {"mini_clip": dict(self.mcfg), "mlsys1": dict(self.syscfg)} + + with open(out_dir / "config.json", "w") as f: + json.dump(payload, f, indent=2) + + def _cuda_mem_stats(self) -> Dict[str, float]: + if self.device != "cuda": + return {} + torch.cuda.synchronize() + alloc = torch.cuda.memory_allocated() / (1024**2) + reserved = torch.cuda.memory_reserved() / (1024**2) + peak = torch.cuda.max_memory_allocated() / (1024**2) + return { + "sys/gpu_mem_alloc_mb": alloc, + "sys/gpu_mem_reserved_mb": reserved, + "sys/gpu_mem_peak_alloc_mb": peak, + } + + def train_one_epoch(self, epoch: int) -> Tuple[float, Dict[str, float]]: + self.model.train() + loader = self.dm.train_dataloader() + meters = StepMeters() + total_loss = 0.0 + + opt_steps = 0 # number of optimizer updates this epoch + grad_accum = int(getattr(self.syscfg, "gradient_accumulation_steps", 1) or 1) + + epoch_start = time.perf_counter() + last_iter_end = time.perf_counter() + + # Autocast context: CUDA only + autocast_ctx = ( + torch.amp.autocast(device_type="cuda", enabled=self.use_amp) + if self.device == "cuda" + else nullcontext() + ) + + for step, (images, text) in enumerate( + tqdm(loader, desc=f"[Train Epoch {epoch}]") + ): + data_t = time.perf_counter() - last_iter_end + step_start = time.perf_counter() + + images = images.to(self.device, non_blocking=True) + text = text.to(self.device, non_blocking=True) + + # Start of an accumulation window + if step % grad_accum == 0: + self.optim.zero_grad(set_to_none=True) + + with autocast_ctx: + img_emb, txt_emb = self.model(images, text) + + # logit_scale in fp32 + clamp for stability (CLIP standard) + logit_scale = self.model.logit_scale.float().clamp(max=4.6052).exp() + + out = self.loss_fn(img_emb, txt_emb, logit_scale) + loss_full = out["loss"] # unscaled (for reporting) + loss = loss_full / grad_accum # scaled (for backward) + + # NaN guard (check the TRUE loss) + if not torch.isfinite(loss_full): + print(f"[NaN] epoch={epoch} step={step} loss={loss_full}") + print("logit_scale:", logit_scale) + raise RuntimeError("Non-finite loss encountered") + + # backward + if self.use_amp: + self.scaler.scale(loss).backward() + else: + loss.backward() + + # Step only every grad_accum steps + if (step + 1) % grad_accum == 0: + if self.use_amp: + self.scaler.step(self.optim) + self.scaler.update() + else: + self.optim.step() + opt_steps += 1 + + if self.device == "cuda": + torch.cuda.synchronize() + + step_t = time.perf_counter() - step_start + bs = images.shape[0] + meters.add(batch_size=bs, step_time_s=step_t, data_time_s=data_t) + + # accumulate TRUE loss for epoch average (not scaled) + total_loss += float(loss_full.item()) + + if self.use_wandb and (step % max(self.mcfg.train.log_every, 1) == 0): + pair_sim = (img_emb * txt_emb).sum(dim=-1) + log_payload = { + "train/loss": float(loss_full.item()), # <- log unscaled loss + "train/loss_img": float(out["loss_img"].item()), + "train/loss_txt": float(out["loss_txt"].item()), + "train/logit_scale": float(logit_scale.item()), + "train/cos_mean": float(pair_sim.mean().item()), + "train/cos_std": float(pair_sim.std().item()), + "train/lr": float(self.optim.param_groups[0]["lr"]), + "sys/steps_per_sec": meters.steps_per_sec, + "sys/samples_per_sec": meters.samples_per_sec, + "sys/data_time_ratio": meters.data_time_ratio, + "sys/grad_accum_steps": int(grad_accum), + "epoch": epoch, + "step": epoch * len(loader) + step, + } + log_payload.update(self._cuda_mem_stats()) + wandb.log(log_payload) + + last_iter_end = time.perf_counter() + + # Flush remaining gradients if dataset size not divisible by grad_accum + if len(loader) % grad_accum != 0: + if self.use_amp: + self.scaler.step(self.optim) + self.scaler.update() + else: + self.optim.step() + opt_steps += 1 + + epoch_time = time.perf_counter() - epoch_start + avg_loss = total_loss / max(len(loader), 1) + + samples_per_sec = meters.samples / max(epoch_time, 1e-9) + updates_per_sec = opt_steps / max(epoch_time, 1e-9) + + epoch_summary = { + "sys/train_epoch_time_s": float(epoch_time), + "sys/train_steps_per_sec_epoch": float( + meters.steps / max(epoch_time, 1e-9) + ), + "sys/train_samples_per_sec_epoch": float(samples_per_sec), + # new metrics + "sys/train_opt_steps_epoch": int(opt_steps), + "sys/train_updates_per_sec_epoch": float(updates_per_sec), + "sys/train_effective_samples_per_sec_epoch": float( + samples_per_sec * grad_accum + ), + "sys/grad_accum_steps": int(grad_accum), + } + epoch_summary.update(self._cuda_mem_stats()) + + return avg_loss, epoch_summary + + def validate(self, epoch: int) -> Tuple[float, float, float, Dict[str, float]]: + self.model.eval() + loader = self.dm.val_dataloader() + total_loss = 0.0 + all_img, all_txt = [], [] + + start = time.perf_counter() + with torch.no_grad(): + for imgs, text in tqdm(loader, desc=f"[Val Epoch {epoch}]"): + imgs = imgs.to(self.device, non_blocking=True) + text = text.to(self.device, non_blocking=True) + + img_emb, txt_emb = self.model(imgs, text) + logit_scale = self.model.logit_scale.exp() + + out = self.loss_fn(img_emb, txt_emb, logit_scale) + total_loss += out["loss"].item() + + all_img.append(img_emb) + all_txt.append(txt_emb) + + if self.device == "cuda": + torch.cuda.synchronize() + + val_time = time.perf_counter() - start + avg_loss = total_loss / max(len(all_img), 1) + + all_img = torch.cat(all_img, dim=0) + all_txt = torch.cat(all_txt, dim=0) + r1_img, r1_txt = compute_recall_at_k(all_img, all_txt, k=1) + + val_summary = {"sys/val_epoch_time_s": val_time} + val_summary.update(self._cuda_mem_stats()) + + if self.use_wandb: + wandb.log( + { + "val/loss": avg_loss, + "val/R@1_img2txt": r1_img, + "val/R@1_txt2img": r1_txt, + "epoch": epoch, + **val_summary, + } + ) + + return avg_loss, r1_img, r1_txt, val_summary + + def fit(self): + warmup_epochs = 10 + total_start = time.perf_counter() + + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + + for epoch in range(self.max_epochs): + train_loss, train_sys = self.train_one_epoch(epoch) + val_loss, r1_i2t, r1_t2i, val_sys = self.validate(epoch) + + self.train_sys_history.append(train_sys) + self.val_sys_history.append(val_sys) + + print( + f"Epoch {epoch:02d} | " + f"train={train_loss:.4f} | " + f"val={val_loss:.4f} | " + f"R@1(img→txt)={r1_i2t:.3f} | " + f"R@1(txt→img)={r1_t2i:.3f} | " + f"train_sps={train_sys.get('sys/train_samples_per_sec_epoch', 0):.1f}" + ) + + # Save best + if val_loss < self.best_val: + self.best_val = val_loss + torch.save(self.model.state_dict(), self.best_ckpt_path) + print(f"✓ Saved BEST checkpoint → {self.best_ckpt_path}") + + # Early stop + if self.early_stop_enabled and epoch >= warmup_epochs: + if val_loss < self.best_val - self.es_min_delta: + self.epochs_no_improve = 0 + else: + self.epochs_no_improve += 1 + if self.epochs_no_improve >= self.es_patience: + print(f"Early stopping at epoch {epoch + 1} (no improvement)") + break + + total_time = time.perf_counter() - total_start + self.dump_run_metrics(total_time) + + if self.use_wandb: + wandb.log({"sys/total_wall_clock_s": total_time}) + wandb.finish() + + +class MiniCLIPTrainerDDPSingleNodeMLsys: + """ + DDP single-node trainer. + Keeps baseline trainer stable by isolating DDP-specific logic here. + """ + + def __init__(self, *, model, datamodule, mcfg, syscfg): + self.model = model + self.dm = datamodule + self.mcfg = mcfg + self.syscfg = syscfg + + # DDP env (torchrun sets these) + self.rank = int(os.environ.get("RANK", "0")) + self.local_rank = int(os.environ.get("LOCAL_RANK", "0")) + self.world_size = int(os.environ.get("WORLD_SIZE", "1")) + self.is_ddp = self.world_size > 1 + self.is_rank0 = self.rank == 0 + + ddp_cfg = getattr(syscfg, "ddp", None) + self.backend = getattr(ddp_cfg, "backend", "nccl") if ddp_cfg else "nccl" + self.find_unused_parameters = bool( + getattr(ddp_cfg, "find_unused_parameters", False) if ddp_cfg else False + ) + + if self.is_rank0: + print( + "DDP env:", + {k: os.environ.get(k) for k in ["RANK", "LOCAL_RANK", "WORLD_SIZE"]}, + ) + + # Device placement (torch.device, not strings) + if torch.cuda.is_available(): + n = torch.cuda.device_count() + if self.local_rank >= n: + raise RuntimeError( + f"LOCAL_RANK={self.local_rank} but only {n} CUDA device(s) visible." + f"Fix: torchrun --nproc_per_node={n} (or set CUDA_VISIBLE_DEVICES)." + ) + torch.cuda.set_device(self.local_rank) + self.device = torch.device("cuda", self.local_rank) + else: + # DDP on CPU is possible with gloo, but your config is nccl, so treat as error + if self.is_ddp and self.backend == "nccl": + raise RuntimeError( + "NCCL backend requires CUDA, but CUDA is not available." + ) + self.device = torch.device("cpu") + + # Init process group (only if multi-process) + if self.is_ddp and not dist.is_initialized(): + dist.init_process_group(backend=self.backend, init_method="env://") + + # Early stopping (rank0 decision; protocol can disable it) + es_cfg = getattr(mcfg.train, "early_stop", None) + self.early_stop_enabled = bool(es_cfg and es_cfg.get("enabled", False)) + self.es_patience = es_cfg.get("patience", 5) if es_cfg else 0 + self.es_min_delta = es_cfg.get("min_delta", 0.0) if es_cfg else 0.0 + self.best_val = float("inf") + self.epochs_no_improve = 0 + + # Protocol (sets max_epochs, seed, can disable early stop) + self._apply_protocol() + + # Move model then wrap DDP + self.model.to(self.device) + if self.is_ddp: + # device_ids/output_device should be ints when using CUDA + self.model = DDP( + self.model, + device_ids=[self.local_rank] if self.device.type == "cuda" else None, + output_device=self.local_rank if self.device.type == "cuda" else None, + find_unused_parameters=self.find_unused_parameters, + ) + + # AMP (optional) + self.use_amp = ( + bool(getattr(syscfg, "amp", False)) and self.device.type == "cuda" + ) + self.scaler = torch.amp.GradScaler("cuda") if self.use_amp else None + + # Optim / loss + self.loss_fn = ClipLoss() + self.optim = AdamW( + self.model.parameters(), + lr=mcfg.train.lr, + weight_decay=mcfg.train.weight_decay, + ) + + # Paths (rank0 only writes) + self.ckpt_dir = ( + Path(mcfg.paths.checkpoints) + / mcfg.exp.experiment_name + / f"mlsys1_{syscfg.name}" + ) + if self.is_rank0: + self.ckpt_dir.mkdir(parents=True, exist_ok=True) + self.best_ckpt_path = self.ckpt_dir / "best_clip.pt" + + # W&B (rank0 only) + self.use_wandb = bool(getattr(syscfg, "use_wandb", False)) and self.is_rank0 + if self.use_wandb: + run_name = ( + f"{mcfg.exp.experiment_name}" + f"__mlsys1-{syscfg.name}" + f"__ddp-ws{self.world_size}" + f"__epochs{self.max_epochs}" + f"__bs{mcfg.train.batch_size}" + ) + wandb.init( + dir=str(mcfg.paths.wandb), + project="mini-clip-mlsys1", + name=run_name, + config={ + "mini_clip": OmegaConf.to_container(mcfg, resolve=True), + "mlsys1": OmegaConf.to_container(syscfg, resolve=True), + "ddp": { + "world_size": self.world_size, + "rank": self.rank, + "local_rank": self.local_rank, + "backend": self.backend, + }, + }, + ) + wandb.watch(self.model, log="gradients", log_freq=200) + + # Sys history (rank0 only writes out) + self.train_sys_history = [] + self.val_sys_history = [] + + # protocol + + def _protocol(self): + return getattr(self.syscfg, "protocol", None) + + def _apply_protocol(self): + p = self._protocol() + if not p or not bool(getattr(p, "enabled", False)): + self.max_epochs = int(self.mcfg.train.epochs) + return + + self.max_epochs = int(getattr(p, "max_epochs", self.mcfg.train.epochs)) + + if bool(getattr(p, "disable_early_stop", False)): + self.early_stop_enabled = False + + seed = getattr(p, "seed", None) + if seed is not None: + seed = int(seed) + int(self.rank) # per-rank offset + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + # helpers + + def _barrier(self): + if self.is_ddp and dist.is_initialized(): + dist.barrier() + + def _reduce_mean(self, x: float) -> float: + if not self.is_ddp: + return float(x) + t = torch.tensor([float(x)], device=self.device) + dist.all_reduce(t, op=dist.ReduceOp.SUM) + t /= float(self.world_size) + return float(t.item()) + + def _cuda_mem_stats(self) -> Dict[str, float]: + if self.device.type != "cuda": + return {} + torch.cuda.synchronize() + alloc = torch.cuda.memory_allocated() / (1024**2) + reserved = torch.cuda.memory_reserved() / (1024**2) + peak = torch.cuda.max_memory_allocated() / (1024**2) + return { + "sys/gpu_mem_alloc_mb": float(alloc), + "sys/gpu_mem_reserved_mb": float(reserved), + "sys/gpu_mem_peak_alloc_mb": float(peak), + } + + # train/val + + def train_one_epoch(self, epoch: int) -> Tuple[float, Dict[str, float]]: + self.model.train() + loader = self.dm.train_dataloader() + + # DistributedSampler epoch seed + sampler = getattr(loader, "sampler", None) + if sampler is not None and hasattr(sampler, "set_epoch"): + sampler.set_epoch(epoch) + + meters = StepMeters() + total_loss = 0.0 + + epoch_start = time.perf_counter() + last_iter_end = time.perf_counter() + # opt_steps = 0 # number of optimizer updates this epoch + + autocast_ctx = ( + torch.amp.autocast(device_type="cuda", enabled=self.use_amp) + if self.device.type == "cuda" + else nullcontext() + ) + + for step, (images, text) in enumerate( + tqdm(loader, desc=f"[Train Epoch {epoch}]") + ): + data_t = time.perf_counter() - last_iter_end + step_start = time.perf_counter() + + images = images.to(self.device, non_blocking=True) + text = text.to(self.device, non_blocking=True) + + with autocast_ctx: + img_emb, txt_emb = self.model(images, text) + + # logit_scale from the underlying module in DDP + base_model = self.model.module if self.is_ddp else self.model + logit_scale = base_model.logit_scale.float().clamp(max=4.6052).exp() + + out = self.loss_fn(img_emb, txt_emb, logit_scale) + loss = out["loss"] + + if not torch.isfinite(loss): + raise RuntimeError( + f"Non-finite loss at epoch={epoch} step={step}: {loss}" + ) + + self.optim.zero_grad(set_to_none=True) + if self.use_amp: + self.scaler.scale(loss).backward() + self.scaler.step(self.optim) + self.scaler.update() + else: + loss.backward() + self.optim.step() + # opt_steps += 1 + + if self.device.type == "cuda": + torch.cuda.synchronize() + + step_t = time.perf_counter() - step_start + bs = int(images.shape[0]) + meters.add(batch_size=bs, step_time_s=step_t, data_time_s=data_t) + + total_loss += float(loss.item()) + last_iter_end = time.perf_counter() + + avg_loss_local = total_loss / max(len(loader), 1) + avg_loss = self._reduce_mean(avg_loss_local) + + epoch_time = time.perf_counter() - epoch_start + + # global throughput estimate + samples_per_sec_local = meters.samples / max(epoch_time, 1e-9) + samples_per_sec_global = samples_per_sec_local * float(self.world_size) + + epoch_summary = { + "sys/train_epoch_time_s": float(epoch_time), + "sys/train_samples_per_sec_epoch": float(samples_per_sec_global), + "sys/train_samples_per_sec_rank": float(samples_per_sec_local), + "sys/world_size": int(self.world_size), + } + epoch_summary.update(self._cuda_mem_stats()) + + return float(avg_loss), epoch_summary + + def validate(self, epoch: int) -> Tuple[float, float, float, Dict[str, float]]: + self.model.eval() + loader = self.dm.val_dataloader() + + total_loss = 0.0 + all_img, all_txt = [], [] + + start = time.perf_counter() + with torch.no_grad(): + for imgs, text in tqdm(loader, desc=f"[Val Epoch {epoch}]"): + imgs = imgs.to(self.device, non_blocking=True) + text = text.to(self.device, non_blocking=True) + + img_emb, txt_emb = self.model(imgs, text) + + base_model = self.model.module if self.is_ddp else self.model + logit_scale = base_model.logit_scale.float().clamp(max=4.6052).exp() + + out = self.loss_fn(img_emb, txt_emb, logit_scale) + total_loss += float(out["loss"].item()) + + # local shard retrieval (global gather can be added later) + all_img.append(img_emb.detach().float().cpu()) + all_txt.append(txt_emb.detach().float().cpu()) + + if self.device.type == "cuda": + torch.cuda.synchronize() + + val_time = time.perf_counter() - start + avg_loss_local = total_loss / max(len(loader), 1) + avg_loss = self._reduce_mean(avg_loss_local) + + all_img = torch.cat(all_img, dim=0) if all_img else torch.empty(0, 1) + all_txt = torch.cat(all_txt, dim=0) if all_txt else torch.empty(0, 1) + r1_img, r1_txt = compute_recall_at_k(all_img, all_txt, k=1) + + val_summary = {"sys/val_epoch_time_s": float(val_time)} + val_summary.update(self._cuda_mem_stats()) + + return float(avg_loss), float(r1_img), float(r1_txt), val_summary + + # dumping + + def dump_config(self): + if not self.is_rank0: + return + out_dir = ( + Path(self.mcfg.paths.outputs) + / self.mcfg.exp.experiment_name + / f"mlsys1_{self.syscfg.name}" + ) + out_dir.mkdir(parents=True, exist_ok=True) + payload = { + "mini_clip": OmegaConf.to_container(self.mcfg, resolve=True), + "mlsys1": OmegaConf.to_container(self.syscfg, resolve=True), + "ddp": { + "world_size": self.world_size, + "rank": self.rank, + "local_rank": self.local_rank, + "backend": self.backend, + }, + } + (out_dir / "config.json").write_text(json.dumps(payload, indent=2)) + + def dump_run_metrics(self, total_time: float): + if not self.is_rank0: + return + + def avg(key, records): + vals = [r.get(key) for r in records if key in r] + return mean(vals) if vals else 0.0 + + metrics = { + "experiment_name": self.mcfg.exp.experiment_name, + "mlsys_name": self.syscfg.name, + "trainer": self.syscfg.trainer, + "batch_size": self.mcfg.train.batch_size, + "num_workers": self.mcfg.train.num_workers, + "epochs_run": len(self.train_sys_history), + "total_wall_clock_s": float(total_time), + "avg_train_samples_per_sec": avg( + "sys/train_samples_per_sec_epoch", self.train_sys_history + ), + "avg_val_epoch_time_s": avg("sys/val_epoch_time_s", self.val_sys_history), + "world_size": int(self.world_size), + "best_val_loss": float(self.best_val), + "gradient_accumulation_steps": int( + getattr(self.syscfg, "gradient_accumulation_steps", 1) or 1 + ), + "effective_batch_size": int(self.mcfg.train.batch_size) + * int(getattr(self.syscfg, "gradient_accumulation_steps", 1) or 1), + "avg_train_updates_per_sec": avg( + "sys/train_updates_per_sec_epoch", self.train_sys_history + ), + "avg_train_opt_steps_per_epoch": avg( + "sys/train_opt_steps_epoch", self.train_sys_history + ), + "avg_train_effective_samples_per_sec": avg( + "sys/train_effective_samples_per_sec_epoch", self.train_sys_history + ), + } + + out_dir = ( + Path(self.mcfg.paths.outputs) + / self.mcfg.exp.experiment_name + / f"mlsys1_{self.syscfg.name}" + ) + out_dir.mkdir(parents=True, exist_ok=True) + + (out_dir / "metrics.json").write_text(json.dumps(metrics, indent=2)) + + hist_file = out_dir / "history.jsonl" + with open(hist_file, "w") as f: + for r in self.train_sys_history: + f.write(json.dumps({"split": "train", **r}) + "\n") + for r in self.val_sys_history: + f.write(json.dumps({"split": "val", **r}) + "\n") + + print(f"[MLSys] Run metrics written to {out_dir / 'metrics.json'}") + print(f"[MLSys] Epoch history written to {hist_file}") + + # main loop + + def fit(self): + if self.is_rank0: + self.dump_config() + + warmup_epochs = 10 + total_start = time.perf_counter() + + if self.device.type == "cuda": + torch.cuda.reset_peak_memory_stats() + + for epoch in range(self.max_epochs): + train_loss, train_sys = self.train_one_epoch(epoch) + val_loss, r1_i2t, r1_t2i, val_sys = self.validate(epoch) + + # rank0 only: track history + printing + checkpoint + early stop + if self.is_rank0: + self.train_sys_history.append(train_sys) + self.val_sys_history.append(val_sys) + + print( + f"Epoch {epoch:02d} | " + f"train={train_loss:.4f} | " + f"val={val_loss:.4f} | " + f"R@1(img→txt)={r1_i2t:.3f} | " + f"R@1(txt→img)={r1_t2i:.3f} | " + f"train_sps={train_sys.get('sys/train_samples_per_sec_epoch', 0):.1f}" + ) + + # Save best + if val_loss < self.best_val: + self.best_val = val_loss + base_model = self.model.module if self.is_ddp else self.model + torch.save(base_model.state_dict(), self.best_ckpt_path) + print(f"✓ Saved BEST checkpoint → {self.best_ckpt_path}") + + # Early stop + if self.early_stop_enabled and epoch >= warmup_epochs: + if val_loss < self.best_val - self.es_min_delta: + self.epochs_no_improve = 0 + else: + self.epochs_no_improve += 1 + if self.epochs_no_improve >= self.es_patience: + print(f"Early stopping at epoch {epoch + 1} (no improvement)") + break + + # keep all ranks aligned each epoch + self._barrier() + + total_time = time.perf_counter() - total_start + if self.is_rank0: + self.dump_run_metrics(total_time) + + if self.use_wandb: + wandb.log({"sys/total_wall_clock_s": float(total_time)}) + wandb.finish() + + self._barrier() + if self.is_ddp and dist.is_initialized(): + dist.destroy_process_group()