From 7cdb2d90c2a8014d255eaf7e87e7ccabd5001e09 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Sun, 23 Aug 2026 22:41:39 +0800 Subject: [PATCH] Add multi-GPU DDP for Omni LoRA training --- README.md | 50 +++++ train_omni_turn_lora.py | 444 +++++++++++++++++++++++++++++++--------- 2 files changed, 402 insertions(+), 92 deletions(-) diff --git a/README.md b/README.md index 856e238..673ac99 100644 --- a/README.md +++ b/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 diff --git a/train_omni_turn_lora.py b/train_omni_turn_lora.py index 981e071..960aed7 100644 --- a/train_omni_turn_lora.py +++ b/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() - print( - f"LoRA matched {len(matched)} language layers; first={matched[0]}, " - f"last={matched[-1]}" - ) + if verbose: + language_model.print_trainable_parameters() + print( + 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: - args = parse_args(argv) +def _run_training(args: argparse.Namespace, run: DistributedRun) -> None: 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", + output_dir / "adapter" / "adapter_config.json", 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) - if device.type != "cuda": - raise ValueError("This LoRA pilot is intended for one CUDA GPU") + device = run.device 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( - f"Cached training examples: {len(dataset)}, hidden_size={dataset.hidden_size}, " - f"batch={args.batch_size} x accum={args.gradient_accumulation_steps}" - ) + if run.is_main: + print( + 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( - args.model_id, - dtype=args.dtype, - device_map=args.device, - use_flash_attention_2=args.flash_attention_2, - ) + base_model, processor = _load_base_model(args, run) 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( - language_model, - inputs_embeds=inputs_embeds, - position_ids=position_ids, - attention_mask=attention_mask, - readout_indices=readout_indices, + synchronize_gradients = ( + accumulation_index == args.gradient_accumulation_steps - 1 ) - logits = head(states.float()) - loss = egospeak_multihorizon_loss(logits, targets) - (loss / args.gradient_accumulation_steps).backward() + with ExitStack() as stack: + if run.enabled and not synchronize_gradients: + 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) - progress.update(1) - progress.set_postfix( - loss=f"{accumulated_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: + history.append(record) + progress.update(1) + progress.set_postfix( + loss=f"{reported_loss:.4f}", + lr=f"{optimizer.param_groups[0]['lr']:.2g}", + ) + 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) - torch.save( - { - "head": head.state_dict(), - "head_config": { - "hidden_size": dataset.hidden_size, - "num_horizons": NUM_HORIZONS, - "num_classes": NUM_CLASSES, - "dropout": head.dropout, + if run.is_main: + raw_language_model.save_pretrained(adapter_dir) + torch.save( + { + "head": raw_head.state_dict(), + "head_config": { + "hidden_size": dataset.hidden_size, + "num_horizons": NUM_HORIZONS, + "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": ( - str(args.head_checkpoint.resolve()) if args.head_checkpoint else None - ), - "training_step": args.max_steps, - }, - output_dir / "turn_head.pt", + output_dir / "turn_head.pt", + ) + global_batch_size = ( + run.world_size * args.batch_size * args.gradient_accumulation_steps ) 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) - print(f"Final cached loss: {final_loss:.6f} (initial {initial_loss:.6f})") - print( - 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: + if run.is_main: + save_json(output_dir / "train_summary.json", summary) + print(f"Final cached loss: {final_loss:.6f} (initial {initial_loss:.6f})") print( - "Online current-tail AV: " - f"{online_timing['online_av_seconds']:.3f}s / " - f"{online_timing['online_av_calls']} calls" + f"Training speed: {training_seconds:.3f}s / {args.max_steps} updates = " + f"{args.max_steps / training_seconds:.3f} update/s; " + 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: " - f"{online_timing['consumer_queue_wait_seconds']:.3f}s; " - f"producer batches={online_timing['producer_batches']}" + "Online current-tail AV: " + f"{online_timing['online_av_seconds']:.3f}s / " + f"{online_timing['online_av_calls']} calls" ) - print(f"Saved LoRA adapter to {adapter_dir}") - print(f"Saved turn head to {output_dir / 'turn_head.pt'}") - print(f"Peak CUDA allocated: {summary['peak_cuda_allocated_gib']:.2f} GiB") - if summary["producer_peak_cuda_allocated_gib"] is not None: + if online_timing["consumer_queue_wait_seconds"] is not None: + print( + "Consumer queue wait: " + 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: " - f"{summary['producer_peak_cuda_allocated_gib']:.2f} GiB" + "Peak CUDA allocated across ranks: " + 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__":