wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add multi-GPU DDP for Omni LoRA training

Browse Source
This commit is contained in:
Xin Wang
2026-08-23 22:41:39 +08:00
parent cb4abfdf82
commit 7cdb2d90c2
2 changed files with 402 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
README.md
View File

@@ -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

334
train_omni_turn_lora.py
View File

@@ -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,6 +688,7 @@ def _load_initial_head(
dropout=float(config["dropout"]),
)
head.load_state_dict(payload.get("head", payload))
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,6 +738,7 @@ def _prepare_language_lora(
)
except TypeError:
language_model.gradient_checkpointing_enable()
if verbose:
language_model.print_trainable_parameters()
print(
f"LoRA matched {len(matched)} language layers; first={matched[0]}, "
@@ -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"
)
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,
)
if run.is_main:
print(
f"Cached training examples: {len(dataset)}, hidden_size={dataset.hidden_size}, "
f"batch={args.batch_size} x accum={args.gradient_accumulation_steps}"
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,
)
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,6 +1123,15 @@ 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)
synchronize_gradients = (
accumulation_index == args.gradient_accumulation_steps - 1
)
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,
@@ -929,34 +1144,38 @@ def main(argv: Sequence[str] | None = None) -> None:
(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"]),
}
if run.is_main:
history.append(record)
progress.update(1)
progress.set_postfix(
loss=f"{accumulated_loss:.4f}",
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,18 +1201,29 @@ 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:
raw_language_model.save_pretrained(adapter_dir)
torch.save(
{
"head": head.state_dict(),
"head": raw_head.state_dict(),
"head_config": {
"hidden_size": dataset.hidden_size,
"num_horizons": NUM_HORIZONS,
"num_classes": NUM_CLASSES,
"dropout": head.dropout,
"dropout": raw_head.dropout,
},
"initial_head_checkpoint": (
str(args.head_checkpoint.resolve()) if args.head_checkpoint else None
@@ -1002,6 +1232,9 @@ def main(argv: Sequence[str] | None = None) -> None:
},
output_dir / "turn_head.pt",
)
global_batch_size = (
run.world_size * args.batch_size * args.gradient_accumulation_steps
)
summary = {
"schema_version": 1,
"status": "pilot" if args.max_examples else "training",
@@ -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,11 +1276,13 @@ def main(argv: Sequence[str] | None = None) -> None:
"args": _jsonable_args(args),
"history": history,
}
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(
f"Training speed: {training_seconds:.3f}s / {args.max_steps} updates = "
f"{args.max_steps / training_seconds:.3f} update/s"
f"{args.max_steps / training_seconds:.3f} update/s; "
f"estimated {summary['estimated_examples_per_second']:.2f} examples/s"
)
if online_timing is not None:
print(
@@ -1052,12 +1298,26 @@ def main(argv: Sequence[str] | None = None) -> None:
)
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")
print(
"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__":