Add multi-GPU DDP for Omni LoRA training

README.md

@@ -970,6 +970,56 @@ manifest contains only selected training anchors, so `train_omni_turn_lora.py`
 uses the sparse dataset automatically and must be run without
 `--online-current-tail`.
 
+### Multi-GPU Thinker-LoRA training
+
+Fully cached AV training supports native single-node DDP through `torchrun`.
+Each rank holds one complete 3B Thinker replica, reads a disjoint shard of the
+selected cache anchors, and synchronizes only the trainable LoRA and turn-head
+gradients.  Model checkpoint loading is deliberately staggered by rank so four
+processes do not read all checkpoint shards simultaneously.  During gradient
+accumulation, DDP synchronization occurs only on the final microbatch.  Only
+rank zero writes the adapter, head, summary, and progress log.
+
+For four GPUs, use per-GPU batch 2 and accumulation 2 for global batch 16:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+OMP_NUM_THREADS=2 \
+MKL_NUM_THREADS=2 \
+OPENBLAS_NUM_THREADS=2 \
+PYTORCH_ALLOC_CONF=expandable_segments:True \
+torchrun --standalone --nproc_per_node=4 train_omni_turn_lora.py \
+  --model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_7 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_11 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_10 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_6 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_9 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_4 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_8 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_5 \
+  --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_sparse_8s/Session_12 \
+  --output-dir checkpoints/omni_turn_lora_easycom_sparse_8s_ddp4 \
+  --max-steps 10170 \
+  --batch-size 2 \
+  --gradient-accumulation-steps 2 \
+  --num-workers 0 \
+  --block-cache-size 32 \
+  --learning-rate 2e-5 \
+  --head-learning-rate 1e-4 \
+  --warmup-steps 500 \
+  --dtype bfloat16 \
+  --flash-attention-2 \
+  --log-every 20
+```
+
+`10170` is only the current three-epoch estimate for roughly 54,000 selected
+anchors.  Recompute it from the completed manifests as
+`ceil(num_examples / 16) * 3`.  Do not pass `--device`: under `torchrun`, each
+process uses its `LOCAL_RANK`.  DDP intentionally rejects
+`--online-current-tail`; the online producer-consumer path remains a separate
+benchmark.
+
 For a speed-only alternative, `train_omni_turn_lora.py` supports a two-GPU
 producer-consumer benchmark.  The producer GPU runs the frozen AV towers and
 places current-tail embeddings in a bounded CPU queue while the consumer GPU

train_omni_turn_lora.py

@@ -16,19 +16,27 @@ from __future__ import annotations
 import argparse
 import gc
 import importlib.metadata
+import inspect
 import json
 import math
+import os
 import queue
 import random
 import threading
 import time
+from contextlib import ExitStack
+from dataclasses import dataclass
+from datetime import timedelta
 from pathlib import Path
 from typing import Any, Sequence
 
 import numpy as np
 import torch
+import torch.distributed as dist
 from torch import nn
+from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
 from tqdm.auto import tqdm
 
 import streaming_turn_model as core
@@ -63,6 +71,130 @@ DEFAULT_LORA_TARGETS = (
 )
 
 
+@dataclass(frozen=True)
+class DistributedRun:
+    enabled: bool
+    rank: int
+    local_rank: int
+    world_size: int
+    device: torch.device
+
+    @property
+    def is_main(self) -> bool:
+        return self.rank == 0
+
+
+def _initialize_distributed(args: argparse.Namespace) -> DistributedRun:
+    """Initialize torchrun/NCCL when WORLD_SIZE is greater than one."""
+
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    if world_size <= 1:
+        device = torch.device(args.device)
+        if device.type != "cuda":
+            raise ValueError("This LoRA trainer requires a CUDA device")
+        torch.cuda.set_device(device)
+        return DistributedRun(
+            enabled=False,
+            rank=0,
+            local_rank=device.index or 0,
+            world_size=1,
+            device=device,
+        )
+
+    if not dist.is_available():
+        raise RuntimeError("torch.distributed is unavailable in this PyTorch build")
+    if "LOCAL_RANK" not in os.environ or "RANK" not in os.environ:
+        raise RuntimeError("Multi-GPU training must be launched with torchrun")
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group(
+        backend="nccl",
+        timeout=timedelta(minutes=float(args.ddp_timeout_minutes)),
+    )
+    return DistributedRun(
+        enabled=True,
+        rank=dist.get_rank(),
+        local_rank=local_rank,
+        world_size=dist.get_world_size(),
+        device=torch.device("cuda", local_rank),
+    )
+
+
+def _distributed_barrier(run: DistributedRun) -> None:
+    if run.enabled:
+        dist.barrier()
+
+
+def _distributed_average(value: float, run: DistributedRun) -> float:
+    if not run.enabled:
+        return float(value)
+    tensor = torch.tensor(value, dtype=torch.float64, device=run.device)
+    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
+    return float((tensor / run.world_size).item())
+
+
+def _synchronize_trainable_parameters(
+    modules: Sequence[torch.nn.Module], run: DistributedRun
+) -> None:
+    """Broadcast only LoRA/head tensors, avoiding a second 3B-model broadcast."""
+
+    if not run.enabled:
+        return
+    for module in modules:
+        for parameter in module.parameters():
+            if parameter.requires_grad:
+                dist.broadcast(parameter.data, src=0)
+
+
+def _wrap_ddp(module: torch.nn.Module, run: DistributedRun) -> torch.nn.Module:
+    if not run.enabled:
+        return module
+    kwargs: dict[str, Any] = {
+        "device_ids": [run.local_rank],
+        "output_device": run.local_rank,
+        "broadcast_buffers": False,
+        "find_unused_parameters": False,
+    }
+    # PyTorch >= 2.7 can skip the expensive initial broadcast of all frozen
+    # base-model parameters.  The checkpoint is identical on every rank and
+    # trainable LoRA/head tensors were synchronized explicitly above.
+    if "init_sync" in inspect.signature(DDP.__init__).parameters:
+        kwargs["init_sync"] = False
+    return DDP(module, **kwargs)
+
+
+def _load_base_model(
+    args: argparse.Namespace, run: DistributedRun
+) -> tuple[torch.nn.Module, Any]:
+    """Load one rank at a time so four shard readers do not hit storage together."""
+
+    if not run.enabled:
+        return core.load_qwen_omni(
+            args.model_id,
+            dtype=args.dtype,
+            device_map=str(run.device),
+            use_flash_attention_2=args.flash_attention_2,
+        )
+
+    loaded: tuple[torch.nn.Module, Any] | None = None
+    for loading_rank in range(run.world_size):
+        if run.rank == loading_rank:
+            print(
+                f"[rank {run.rank}/{run.world_size}] loading model on {run.device}",
+                flush=True,
+            )
+            loaded = core.load_qwen_omni(
+                args.model_id,
+                dtype=args.dtype,
+                device_map=str(run.device),
+                use_flash_attention_2=args.flash_attention_2,
+            )
+        dist.barrier()
+    if loaded is None:
+        raise RuntimeError(f"Rank {run.rank} did not load a model")
+    return loaded
+
+
 def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--model-id", required=True)
@@ -162,7 +294,17 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
         choices=("bfloat16", "float16"),
         default="bfloat16",
     )
-    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument(
+        "--device",
+        default="cuda:0",
+        help="Single-GPU device; torchrun overrides this with LOCAL_RANK",
+    )
+    parser.add_argument(
+        "--ddp-timeout-minutes",
+        type=float,
+        default=60.0,
+        help="NCCL process-group timeout used by torchrun multi-GPU training",
+    )
     parser.add_argument("--flash-attention-2", action="store_true")
     parser.add_argument("--overwrite-output", action="store_true")
     parser.add_argument("--no-progress", action="store_true")
@@ -524,6 +666,7 @@ def _load_initial_head(
     *,
     hidden_size: int,
     dropout: float,
+    verbose: bool = True,
 ) -> TurnPredictionHead:
     if checkpoint is None:
         return TurnPredictionHead(hidden_size=hidden_size, dropout=dropout)
@@ -545,7 +688,8 @@ def _load_initial_head(
         dropout=float(config["dropout"]),
     )
     head.load_state_dict(payload.get("head", payload))
-    print(f"Initialized classification head from {checkpoint}")
+    if verbose:
+        print(f"Initialized classification head from {checkpoint}")
     return head
 
 
@@ -557,6 +701,7 @@ def _prepare_language_lora(
     dropout: float,
     target_modules: Sequence[str],
     gradient_checkpointing: bool,
+    verbose: bool = True,
 ) -> torch.nn.Module:
     try:
         from peft import LoraConfig, get_peft_model
@@ -593,11 +738,12 @@ def _prepare_language_lora(
             )
         except TypeError:
             language_model.gradient_checkpointing_enable()
-    language_model.print_trainable_parameters()
+    if verbose:
-    print(
+        language_model.print_trainable_parameters()
-        f"LoRA matched {len(matched)} language layers; first={matched[0]}, "
+        print(
-        f"last={matched[-1]}"
+            f"LoRA matched {len(matched)} language layers; first={matched[0]}, "
-    )
+            f"last={matched[-1]}"
+        )
     return language_model
 
 
@@ -618,13 +764,15 @@ def _evaluate_loss(
     device: torch.device,
     model_dtype: torch.dtype,
     max_batches: int,
+    distributed_run: DistributedRun,
     materialize_batch: Any | None = None,
 ) -> float:
     language_was_training = language_model.training
     head_was_training = head.training
     language_model.eval()
     head.eval()
-    losses: list[float] = []
+    loss_sum = 0.0
+    example_count = 0
     for batch_index, batch in enumerate(loader):
         if max_batches > 0 and batch_index >= max_batches:
             break
@@ -645,10 +793,19 @@ def _evaluate_loss(
             readout_indices=readout_indices,
         )
         logits = head(states.float())
-        losses.append(float(egospeak_multihorizon_loss(logits, targets).item()))
+        batch_size = int(targets.shape[0])
+        loss_sum += float(egospeak_multihorizon_loss(logits, targets).item()) * batch_size
+        example_count += batch_size
     language_model.train(language_was_training)
     head.train(head_was_training)
-    return float(np.mean(losses))
+    totals = torch.tensor(
+        [loss_sum, float(example_count)], dtype=torch.float64, device=device
+    )
+    if distributed_run.enabled:
+        dist.all_reduce(totals, op=dist.ReduceOp.SUM)
+    if float(totals[1].item()) <= 0:
+        raise RuntimeError("Evaluation loader produced no examples")
+    return float((totals[0] / totals[1]).item())
 
 
 def _jsonable_args(args: argparse.Namespace) -> dict[str, Any]:
@@ -663,13 +820,17 @@ def _jsonable_args(args: argparse.Namespace) -> dict[str, Any]:
     return output
 
 
-def main(argv: Sequence[str] | None = None) -> None:
+def _run_training(args: argparse.Namespace, run: DistributedRun) -> None:
-    args = parse_args(argv)
     if args.max_steps <= 0:
         raise ValueError("--max-steps must be positive")
     if args.batch_size <= 0 or args.gradient_accumulation_steps <= 0:
         raise ValueError("Batch size and gradient accumulation must be positive")
     if args.online_current_tail:
+        if run.enabled:
+            raise ValueError(
+                "DDP currently supports fully cached AV training only; remove "
+                "--online-current-tail when launching with torchrun"
+            )
         if args.easycom_session_dir is None:
             raise ValueError(
                 "--easycom-session-dir is required by --online-current-tail"
@@ -689,6 +850,8 @@ def main(argv: Sequence[str] | None = None) -> None:
             raise ValueError("--producer-device must be a CUDA device such as cuda:1")
     if args.prefetch_depth <= 0:
         raise ValueError("--prefetch-depth must be positive")
+    if args.ddp_timeout_minutes <= 0:
+        raise ValueError("--ddp-timeout-minutes must be positive")
     if (
         args.audio_right_padding_samples is not None
         and args.audio_right_padding_samples < 0
@@ -701,15 +864,14 @@ def main(argv: Sequence[str] | None = None) -> None:
         raise ValueError(
             "Audio preprocessing overrides apply only with --online-current-tail"
         )
-    if args.num_workers != 0:
+    if args.num_workers != 0 and run.is_main:
         print(
             "Warning: each worker owns an independent block-file LRU; use 0 for "
             "the 16-example pilot and increase only after profiling."
         )
     output_dir = args.output_dir.expanduser().resolve()
     known_outputs = (
-        output_dir / "adapter_config.json",
+        output_dir / "adapter" / "adapter_config.json",
-        output_dir / "adapter",
         output_dir / "turn_head.pt",
         output_dir / "train_summary.json",
     )
@@ -717,12 +879,12 @@ def main(argv: Sequence[str] | None = None) -> None:
         raise FileExistsError(
             f"Training output already exists in {output_dir}; pass --overwrite-output"
         )
-    output_dir.mkdir(parents=True, exist_ok=True)
+    if run.is_main:
+        output_dir.mkdir(parents=True, exist_ok=True)
+    _distributed_barrier(run)
     _seed_everything(args.seed)
     core._configure_qwen_logging()
-    device = torch.device(args.device)
+    device = run.device
-    if device.type != "cuda":
-        raise ValueError("This LoRA pilot is intended for one CUDA GPU")
 
     dataset = OmniTurnCacheDataset(
         args.cache_dir,
@@ -737,11 +899,30 @@ def main(argv: Sequence[str] | None = None) -> None:
         if args.online_current_tail
         else collate_cached_sequences
     )
-    generator = torch.Generator().manual_seed(args.seed)
+    generator = torch.Generator().manual_seed(args.seed + run.rank)
+    train_sampler: DistributedSampler | None = None
+    eval_sampler: DistributedSampler | None = None
+    if run.enabled:
+        train_sampler = DistributedSampler(
+            dataset,
+            num_replicas=run.world_size,
+            rank=run.rank,
+            shuffle=True,
+            seed=args.seed,
+            drop_last=False,
+        )
+        eval_sampler = DistributedSampler(
+            dataset,
+            num_replicas=run.world_size,
+            rank=run.rank,
+            shuffle=False,
+            drop_last=False,
+        )
     train_loader = DataLoader(
         dataset,
         batch_size=args.batch_size,
-        shuffle=True,
+        shuffle=train_sampler is None,
+        sampler=train_sampler,
         num_workers=args.num_workers,
         pin_memory=True,
         persistent_workers=args.num_workers > 0,
@@ -752,21 +933,21 @@ def main(argv: Sequence[str] | None = None) -> None:
         dataset,
         batch_size=args.batch_size,
         shuffle=False,
+        sampler=eval_sampler,
         num_workers=0,
         pin_memory=True,
         collate_fn=collate_fn,
     )
-    print(
+    if run.is_main:
-        f"Cached training examples: {len(dataset)}, hidden_size={dataset.hidden_size}, "
+        print(
-        f"batch={args.batch_size} x accum={args.gradient_accumulation_steps}"
+            f"Cached training examples: {len(dataset)}, hidden_size={dataset.hidden_size}, "
-    )
+            f"world={run.world_size}, per_gpu_batch={args.batch_size}, "
+            f"accum={args.gradient_accumulation_steps}, "
+            "global_batch="
+            f"{run.world_size * args.batch_size * args.gradient_accumulation_steps}"
+        )
 
-    base_model, processor = core.load_qwen_omni(
+    base_model, processor = _load_base_model(args, run)
-        args.model_id,
-        dtype=args.dtype,
-        device_map=args.device,
-        use_flash_attention_2=args.flash_attention_2,
-    )
     target_modules = tuple(
         value.strip() for value in args.lora_target_modules.split(",") if value.strip()
     )
@@ -779,6 +960,7 @@ def main(argv: Sequence[str] | None = None) -> None:
         dropout=args.lora_dropout,
         target_modules=target_modules,
         gradient_checkpointing=args.gradient_checkpointing,
+        verbose=run.is_main,
     )
     materialize_batch: OnlineCurrentTailMaterializer | None = None
     producer_model: torch.nn.Module | None = None
@@ -825,7 +1007,7 @@ def main(argv: Sequence[str] | None = None) -> None:
         if args.producer_device is not None:
             print(
                 f"Producer-consumer enabled: frozen AV={args.producer_device}, "
-                f"Thinker LoRA={args.device}, queue={args.prefetch_depth}"
+                f"Thinker LoRA={device}, queue={args.prefetch_depth}"
             )
             # The consumer retains the same language object through its PEFT
             # wrapper, so the outer multimodal container and towers can go.
@@ -844,15 +1026,29 @@ def main(argv: Sequence[str] | None = None) -> None:
         args.head_checkpoint,
         hidden_size=dataset.hidden_size,
         dropout=args.head_dropout,
+        verbose=run.is_main,
     ).to(device=device, dtype=torch.float32)
 
-    lora_parameters = [parameter for parameter in language_model.parameters() if parameter.requires_grad]
+    raw_language_model = language_model
+    raw_head = head
+    _synchronize_trainable_parameters((raw_language_model, raw_head), run)
+    language_model = _wrap_ddp(raw_language_model, run)
+    head = _wrap_ddp(raw_head, run)
+    # Model initialization is synchronized above; rank-specific RNG streams
+    # avoid identical dropout masks for different distributed samples.
+    _seed_everything(args.seed + run.rank)
+
+    lora_parameters = [
+        parameter
+        for parameter in raw_language_model.parameters()
+        if parameter.requires_grad
+    ]
     if not lora_parameters:
         raise RuntimeError("LoRA preparation produced no trainable parameters")
     optimizer = torch.optim.AdamW(
         [
             {"params": lora_parameters, "lr": args.learning_rate},
-            {"params": head.parameters(), "lr": args.head_learning_rate},
+            {"params": raw_head.parameters(), "lr": args.head_learning_rate},
         ],
         weight_decay=args.weight_decay,
     )
@@ -870,14 +1066,18 @@ def main(argv: Sequence[str] | None = None) -> None:
         device=device,
         model_dtype=model_dtype,
         max_batches=args.eval_max_batches,
+        distributed_run=run,
         materialize_batch=materialize_batch,
     )
-    print(f"Initial cached loss: {initial_loss:.6f}")
+    if run.is_main:
+        print(f"Initial cached loss: {initial_loss:.6f}")
     if materialize_batch is not None:
         materialize_batch.reset_timing()
     language_model.train()
     head.train()
     optimizer.zero_grad(set_to_none=True)
+    _distributed_barrier(run)
+    torch.cuda.synchronize(device)
     training_started = time.perf_counter()
     batch_producer: AsyncTailBatchProducer | None = None
     if materialize_batch is not None and args.producer_device is not None:
@@ -888,17 +1088,20 @@ def main(argv: Sequence[str] | None = None) -> None:
             prefetch_depth=args.prefetch_depth,
         )
         batch_producer.start()
+    sampler_epoch = 0
+    if train_sampler is not None:
+        train_sampler.set_epoch(sampler_epoch)
     iterator = iter(train_loader) if batch_producer is None else None
     history: list[dict[str, float | int]] = []
     progress = tqdm(
         total=args.max_steps,
         desc="Thinker LoRA",
         unit="update",
-        disable=args.no_progress,
+        disable=args.no_progress or not run.is_main,
     )
     for update_step in range(1, args.max_steps + 1):
         accumulated_loss = 0.0
-        for _ in range(args.gradient_accumulation_steps):
+        for accumulation_index in range(args.gradient_accumulation_steps):
             if batch_producer is not None:
                 batch = batch_producer.next_batch()
             else:
@@ -906,6 +1109,9 @@ def main(argv: Sequence[str] | None = None) -> None:
                 try:
                     batch = next(iterator)
                 except StopIteration:
+                    sampler_epoch += 1
+                    if train_sampler is not None:
+                        train_sampler.set_epoch(sampler_epoch)
                     iterator = iter(train_loader)
                     batch = next(iterator)
             if materialize_batch is not None and batch_producer is None:
@@ -917,46 +1123,59 @@ def main(argv: Sequence[str] | None = None) -> None:
             attention_mask = batch["attention_mask"].to(device=device, non_blocking=True)
             readout_indices = batch["readout_indices"].to(device=device, non_blocking=True)
             targets = batch["targets"].to(device=device, dtype=torch.float32)
-            states = forward_last_state(
+            synchronize_gradients = (
-                language_model,
+                accumulation_index == args.gradient_accumulation_steps - 1
-                inputs_embeds=inputs_embeds,
-                position_ids=position_ids,
-                attention_mask=attention_mask,
-                readout_indices=readout_indices,
             )
-            logits = head(states.float())
+            with ExitStack() as stack:
-            loss = egospeak_multihorizon_loss(logits, targets)
+                if run.enabled and not synchronize_gradients:
-            (loss / args.gradient_accumulation_steps).backward()
+                    assert isinstance(language_model, DDP)
+                    assert isinstance(head, DDP)
+                    stack.enter_context(language_model.no_sync())
+                    stack.enter_context(head.no_sync())
+                states = forward_last_state(
+                    language_model,
+                    inputs_embeds=inputs_embeds,
+                    position_ids=position_ids,
+                    attention_mask=attention_mask,
+                    readout_indices=readout_indices,
+                )
+                logits = head(states.float())
+                loss = egospeak_multihorizon_loss(logits, targets)
+                (loss / args.gradient_accumulation_steps).backward()
             accumulated_loss += float(loss.detach()) / args.gradient_accumulation_steps
 
-        trainable_parameters = [*lora_parameters, *head.parameters()]
+        trainable_parameters = [*lora_parameters, *raw_head.parameters()]
         grad_norm = torch.nn.utils.clip_grad_norm_(
             trainable_parameters, args.max_grad_norm
         )
         optimizer.step()
         scheduler.step()
         optimizer.zero_grad(set_to_none=True)
+        reported_loss = _distributed_average(accumulated_loss, run)
         record = {
             "step": update_step,
-            "loss": accumulated_loss,
+            "loss": reported_loss,
             "grad_norm": float(grad_norm),
             "lora_lr": float(optimizer.param_groups[0]["lr"]),
             "head_lr": float(optimizer.param_groups[1]["lr"]),
         }
-        history.append(record)
+        if run.is_main:
-        progress.update(1)
+            history.append(record)
-        progress.set_postfix(
+            progress.update(1)
-            loss=f"{accumulated_loss:.4f}",
+            progress.set_postfix(
-            lr=f"{optimizer.param_groups[0]['lr']:.2g}",
+                loss=f"{reported_loss:.4f}",
-        )
+                lr=f"{optimizer.param_groups[0]['lr']:.2g}",
-        if update_step == 1 or update_step % args.log_every == 0:
+            )
+        if run.is_main and (update_step == 1 or update_step % args.log_every == 0):
             print(
-                f"step={update_step} loss={accumulated_loss:.6f} "
+                f"step={update_step} loss={reported_loss:.6f} "
                 f"grad_norm={float(grad_norm):.4f} "
                 f"lora_lr={optimizer.param_groups[0]['lr']:.3g}",
                 flush=True,
             )
     progress.close()
+    torch.cuda.synchronize(device)
+    _distributed_barrier(run)
     training_seconds = time.perf_counter() - training_started
     if batch_producer is not None:
         batch_producer.close()
@@ -982,25 +1201,39 @@ def main(argv: Sequence[str] | None = None) -> None:
         device=device,
         model_dtype=model_dtype,
         max_batches=args.eval_max_batches,
+        distributed_run=run,
         materialize_batch=materialize_batch,
     )
+    local_peak_gib = torch.cuda.max_memory_allocated(device) / (1024 ** 3)
+    peak_tensor = torch.tensor(local_peak_gib, dtype=torch.float64, device=device)
+    if run.enabled:
+        gathered_peaks = [torch.zeros_like(peak_tensor) for _ in range(run.world_size)]
+        dist.all_gather(gathered_peaks, peak_tensor)
+        per_rank_peak_gib = [float(value.item()) for value in gathered_peaks]
+    else:
+        per_rank_peak_gib = [local_peak_gib]
+
     adapter_dir = output_dir / "adapter"
-    language_model.save_pretrained(adapter_dir)
+    if run.is_main:
-    torch.save(
+        raw_language_model.save_pretrained(adapter_dir)
-        {
+        torch.save(
-            "head": head.state_dict(),
+            {
-            "head_config": {
+                "head": raw_head.state_dict(),
-                "hidden_size": dataset.hidden_size,
+                "head_config": {
-                "num_horizons": NUM_HORIZONS,
+                    "hidden_size": dataset.hidden_size,
-                "num_classes": NUM_CLASSES,
+                    "num_horizons": NUM_HORIZONS,
-                "dropout": head.dropout,
+                    "num_classes": NUM_CLASSES,
+                    "dropout": raw_head.dropout,
+                },
+                "initial_head_checkpoint": (
+                    str(args.head_checkpoint.resolve()) if args.head_checkpoint else None
+                ),
+                "training_step": args.max_steps,
             },
-            "initial_head_checkpoint": (
+            output_dir / "turn_head.pt",
-                str(args.head_checkpoint.resolve()) if args.head_checkpoint else None
+        )
-            ),
+    global_batch_size = (
-            "training_step": args.max_steps,
+        run.world_size * args.batch_size * args.gradient_accumulation_steps
-        },
-        output_dir / "turn_head.pt",
     )
     summary = {
         "schema_version": 1,
@@ -1012,11 +1245,22 @@ def main(argv: Sequence[str] | None = None) -> None:
         "final_loss": final_loss,
         "training_seconds": training_seconds,
         "updates_per_second": args.max_steps / training_seconds,
+        "estimated_examples_per_second": (
+            args.max_steps * global_batch_size / training_seconds
+        ),
+        "distributed": {
+            "enabled": run.enabled,
+            "world_size": run.world_size,
+            "per_gpu_batch_size": args.batch_size,
+            "gradient_accumulation_steps": args.gradient_accumulation_steps,
+            "global_batch_size": global_batch_size,
+            "per_rank_peak_cuda_allocated_gib": per_rank_peak_gib,
+        },
         "online_tail_timing": online_timing,
         "loss_reduction_ratio": (
             (initial_loss - final_loss) / initial_loss if initial_loss else None
         ),
-        "peak_cuda_allocated_gib": torch.cuda.max_memory_allocated(device) / (1024 ** 3),
+        "peak_cuda_allocated_gib": max(per_rank_peak_gib),
         "producer_peak_cuda_allocated_gib": (
             torch.cuda.max_memory_allocated(torch.device(args.producer_device))
             / (1024 ** 3)
@@ -1032,32 +1276,48 @@ def main(argv: Sequence[str] | None = None) -> None:
         "args": _jsonable_args(args),
         "history": history,
     }
-    save_json(output_dir / "train_summary.json", summary)
+    if run.is_main:
-    print(f"Final cached loss: {final_loss:.6f} (initial {initial_loss:.6f})")
+        save_json(output_dir / "train_summary.json", summary)
-    print(
+        print(f"Final cached loss: {final_loss:.6f} (initial {initial_loss:.6f})")
-        f"Training speed: {training_seconds:.3f}s / {args.max_steps} updates = "
-        f"{args.max_steps / training_seconds:.3f} update/s"
-    )
-    if online_timing is not None:
         print(
-            "Online current-tail AV: "
+            f"Training speed: {training_seconds:.3f}s / {args.max_steps} updates = "
-            f"{online_timing['online_av_seconds']:.3f}s / "
+            f"{args.max_steps / training_seconds:.3f} update/s; "
-            f"{online_timing['online_av_calls']} calls"
+            f"estimated {summary['estimated_examples_per_second']:.2f} examples/s"
         )
-        if online_timing["consumer_queue_wait_seconds"] is not None:
+        if online_timing is not None:
             print(
-                "Consumer queue wait: "
+                "Online current-tail AV: "
-                f"{online_timing['consumer_queue_wait_seconds']:.3f}s; "
+                f"{online_timing['online_av_seconds']:.3f}s / "
-                f"producer batches={online_timing['producer_batches']}"
+                f"{online_timing['online_av_calls']} calls"
             )
-    print(f"Saved LoRA adapter to {adapter_dir}")
+            if online_timing["consumer_queue_wait_seconds"] is not None:
-    print(f"Saved turn head to {output_dir / 'turn_head.pt'}")
+                print(
-    print(f"Peak CUDA allocated: {summary['peak_cuda_allocated_gib']:.2f} GiB")
+                    "Consumer queue wait: "
-    if summary["producer_peak_cuda_allocated_gib"] is not None:
+                    f"{online_timing['consumer_queue_wait_seconds']:.3f}s; "
+                    f"producer batches={online_timing['producer_batches']}"
+                )
+        print(f"Saved LoRA adapter to {adapter_dir}")
+        print(f"Saved turn head to {output_dir / 'turn_head.pt'}")
         print(
-            "Producer peak CUDA allocated: "
+            "Peak CUDA allocated across ranks: "
-            f"{summary['producer_peak_cuda_allocated_gib']:.2f} GiB"
+            f"{summary['peak_cuda_allocated_gib']:.2f} GiB"
         )
+        if summary["producer_peak_cuda_allocated_gib"] is not None:
+            print(
+                "Producer peak CUDA allocated: "
+                f"{summary['producer_peak_cuda_allocated_gib']:.2f} GiB"
+            )
+    _distributed_barrier(run)
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    args = parse_args(argv)
+    run = _initialize_distributed(args)
+    try:
+        _run_training(args, run)
+    finally:
+        if run.enabled and dist.is_initialized():
+            dist.destroy_process_group()
 
 
 if __name__ == "__main__":