# src/trainer.py """ FSDP trainer for CodonTranslator. No frameworks, no sugar. The model computes its own loss. Batch invariants: - codon_ids [B, T] (right-padded; EOS already in-sequence) - species_ids [B] (SpeciesEmbeddingStore provides fixed-size or sequence embeddings) - protein_seqs: list[str] (ESM tokenization happens inside the model) Rules: - If your loader is IterableDataset, you MUST set args.max_steps > 0. We don't guess. - If you want epoch-based, use a sized dataset; we call len(dataloader). """ from __future__ import annotations import os import json import math import re import shutil import logging import time from dataclasses import dataclass import datetime import warnings import importlib.util import inspect from typing import Any, Callable, Dict, Optional, Tuple, List from tqdm import tqdm import torch import torch.nn as nn import torch.distributed as dist from torch.utils.data import DataLoader, IterableDataset from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import ( ShardingStrategy, MixedPrecision, StateDictType, FullStateDictConfig, FullOptimStateDictConfig, ) from safetensors.torch import save_file, load_file import wandb logger = logging.getLogger(__name__) # ------------------------------ # Args # ------------------------------ @dataclass class TrainingArguments: # Output output_dir: str = "checkpoints" save_steps: int = 1000 save_total_limit: int = 3 save_safetensors: bool = True ckpt_recent_window_steps: int = 0 ckpt_recent_interval: int = 0 ckpt_archive_interval: int = 0 # Schedule num_train_epochs: int = 1 max_steps: int = -1 # required for IterableDataset gradient_accumulation_steps: int = 1 warmup_ratio: float = 0.0 lr_scheduler_type: str = "cosine" # "linear" | "cosine" | "constant" # For streaming datasets: if max_steps<0 and steps_per_epoch>0, shape schedule using # total_steps = num_train_epochs * steps_per_epoch steps_per_epoch: int = 0 # Optim learning_rate: float = 5e-4 weight_decay: float = 0.0 adam_beta1: float = 0.9 adam_beta2: float = 0.95 max_grad_norm: float = 1.0 # Data per_device_train_batch_size: int = 8 per_device_eval_batch_size: int = 8 dataloader_num_workers: int = 0 # Precision / dist fp16: bool = False bf16: bool = False fsdp: Optional[str] = None # "full_shard" or None gradient_checkpointing: bool = False # Global hard cap (prefix + start + codon) max_length: int = 4096 # ESM (metadata only; model owns ESM) esm_model_name: str = "esmc_300m" esm_device: str = "cuda" esm_dtype: str = "bf16" # Logging / eval logging_steps: int = 100 eval_steps: int = 0 # streaming eval: limit number of eval batches when eval dataset is Iterable eval_interval: int = 0 # run evaluation every N optimizer steps (0 disables) override_lr_on_resume: bool = False # Minimal data stream resume cursor (stores total samples yielded so far for train dataset). # When provided, we load 'skip_samples' from this JSON at start and set the dataset # to skip exactly that many samples on resume. We also update the file in _save_checkpoint(). data_cursor_path: Optional[str] = None # ------------------------------ # Trainer # ------------------------------ class Trainer: def __init__( self, model: nn.Module, args: TrainingArguments, data_collator: Optional[Callable] = None, train_dataset: Optional[Any] = None, eval_dataset: Optional[Any] = None, tokenizer: Optional[Any] = None, model_init: Optional[Callable[[], nn.Module]] = None, compute_metrics: Optional[Callable] = None, callbacks: Optional[list] = None, optimizers: Tuple[Optional[torch.optim.Optimizer], Optional[Any]] = (None, None), preprocess_logits_for_metrics: Optional[Callable] = None, species_store=None, resume_from_checkpoint: Optional[str] = None, ): self.model = model self.args = args self.tokenizer = tokenizer self.optimizer = optimizers[0] self.lr_scheduler = optimizers[1] self.species_store = species_store self.train_dataloader: Optional[DataLoader] = None self.eval_dataloader: Optional[DataLoader] = None # Device (robust local rank resolution) self.local_rank = 0 if torch.cuda.is_available(): lr_env = os.environ.get("LOCAL_RANK") if lr_env is not None: self.local_rank = int(lr_env) else: r = int(os.environ.get("RANK", "0")) ng = max(1, torch.cuda.device_count()) self.local_rank = (r % ng) self.device = torch.device(f"cuda:{self.local_rank}") torch.cuda.set_device(self.device) cd = torch.cuda.current_device() nm = torch.cuda.get_device_name(cd) logger.info( f"[dist] RANK={os.environ.get('RANK')} LOCAL_RANK={os.environ.get('LOCAL_RANK')} WORLD_SIZE={os.environ.get('WORLD_SIZE')} " f"cuda.count={torch.cuda.device_count()} select={self.device} current={cd} name={nm}" ) else: self.device = torch.device("cpu") # Gradient checkpointing toggle (model owns the flag) base = self._unwrap(self.model) if self.args.gradient_checkpointing and hasattr(base, "gradient_checkpointing"): base.gradient_checkpointing = True # FSDP or single GPU if self.args.fsdp: self._setup_fsdp() else: self.model = self.model.to(self.device) # AMP setup (use torch.amp APIs; GradScaler on CUDA only) self._use_amp = (self.device.type == "cuda") and (self.args.fp16 or self.args.bf16) self._amp_dtype = torch.float16 if self.args.fp16 else (torch.bfloat16 if self.args.bf16 else None) use_cuda = (self.device.type == "cuda") self._scaler = torch.amp.GradScaler(device="cuda", enabled=(use_cuda and self.args.fp16)) self.state = {"epoch": 0, "global_step": 0} # Defer resume until after dataloaders are attached so scheduler can be shaped. self._resume_path: Optional[str] = resume_from_checkpoint # ---- dataloaders ---- def attach_dataloaders(self, train_loader: DataLoader, eval_loader: Optional[DataLoader] = None): # Your dataset should handle sharding. We don't wrap with DistributedSampler here. self.train_dataloader = train_loader self.eval_dataloader = eval_loader # Apply minimal resume cursor to the training dataset if configured p = getattr(self.args, "data_cursor_path", None) if p and os.path.exists(p): with open(p, "r") as f: js = json.load(f) ds = getattr(self.train_dataloader, "dataset", None) if hasattr(ds, "set_resume_skip"): distributed = dist.is_available() and dist.is_initialized() world = dist.get_world_size() if distributed else 1 rank = dist.get_rank() if distributed else 0 # Prefer the total cursor and split evenly across current world size. # If total is missing, sum any saved per_rank list. total: int = 0 if isinstance(js, dict): try: total = int(js.get("skip_samples", 0) or 0) except Exception: total = 0 if total <= 0: raw = js.get("per_rank") if isinstance(raw, list) and raw: try: total = int(sum(int(x) for x in raw)) except Exception: total = 0 if total > 0: if distributed: per = total // max(world, 1) rem = total % max(world, 1) n_rank = per + (1 if rank < rem else 0) ds.set_resume_skip(int(n_rank)) if self._is_main(): logger.info( "resume cursor: total=%s split across world=%s → rank=%s skip=%s", total, world, rank, n_rank, ) else: ds.set_resume_skip(int(total)) if self._is_main(): logger.info("resume cursor: total=%s (single-process) skip=%s", total, total) # ---- optim + scheduler ---- def _create_optimizer_and_scheduler(self): if self.optimizer is None: decay, no_decay = [], [] for n, p in self._unwrap(self.model).named_parameters(): if not p.requires_grad: continue if n.endswith("bias") or "norm" in n.lower() or "ln_" in n.lower(): no_decay.append(p) else: decay.append(p) opt_kwargs = dict( lr=self.args.learning_rate, betas=(self.args.adam_beta1, self.args.adam_beta2), ) params = [ {"params": decay, "weight_decay": self.args.weight_decay}, {"params": no_decay, "weight_decay": 0.0}, ] sig_adamw = inspect.signature(torch.optim.AdamW) if torch.cuda.is_available() and "fused" in sig_adamw.parameters: opt_kwargs["fused"] = True # type: ignore[assignment] self.optimizer = torch.optim.AdamW(params, **opt_kwargs) # Report fused/foreach settings (rank0 only) if self._is_main(): fused_flag = None foreach_flag = None if hasattr(self.optimizer, "defaults"): fused_flag = self.optimizer.defaults.get("fused") foreach_flag = self.optimizer.defaults.get("foreach") logger.info(f"AdamW configured: fused={fused_flag} foreach={foreach_flag}") # total steps and schedule shape ds = getattr(self.train_dataloader, "dataset", None) ga = max(1, self.args.gradient_accumulation_steps) if isinstance(ds, IterableDataset): if self.args.max_steps > 0: # Use max_steps to shape the scheduler; allow multiple epochs to re-iterate the stream steps_per_epoch = self.args.max_steps total_steps = self.args.max_steps elif getattr(self.args, "steps_per_epoch", 0) and self.args.steps_per_epoch > 0: # steps_per_epoch is already expressed in optimizer steps (train.py accounts for grad_accum) steps_per_epoch = max(1, int(self.args.steps_per_epoch)) total_steps = max(1, self.args.num_train_epochs) * steps_per_epoch else: # Unknown epoch size; use constant LR without pre-shaped schedule self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lambda step: 1.0) return else: # sized dataloader: len(dataloader) is number of batches steps_per_epoch = max(len(self.train_dataloader) // ga, 1) total_steps = self.args.max_steps if self.args.max_steps > 0 else self.args.num_train_epochs * steps_per_epoch warmup = int(self.args.warmup_ratio * total_steps) if self.args.lr_scheduler_type == "constant": self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lambda step: 1.0) return def lrs_lambda(step: int) -> float: if step < warmup: return max(float(step) / max(warmup, 1), 1e-6) t = (step - warmup) / max(total_steps - warmup, 1) if self.args.lr_scheduler_type == "linear": return max(1.0 - t, 0.0) # cosine default return 0.5 * (1.0 + math.cos(math.pi * t)) self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lrs_lambda) # ---- training ---- def train(self) -> Dict[str, float]: assert self.train_dataloader is not None, "Call attach_dataloaders() first" # If a resume path was provided, load it now (dataloaders are attached). if getattr(self, "_resume_path", None): self._resume_from(self._resume_path) # loads model/optimizer/scheduler/state self._resume_path = None if self.optimizer is None: self._create_optimizer_and_scheduler() ds = self.train_dataloader.dataset # Exact step budget for streaming datasets when max_steps<0 and steps_per_epoch>0 target_total_steps: Optional[int] = None if isinstance(ds, IterableDataset) and int(self.args.max_steps) < 0: spe = int(getattr(self.args, "steps_per_epoch", 0) or 0) if spe > 0: target_total_steps = max(1, int(self.args.num_train_epochs)) * spe # Determine total steps for progress bar progress_total: Optional[int] = None if int(self.args.max_steps) > 0: progress_total = int(self.args.max_steps) elif isinstance(ds, IterableDataset): if target_total_steps is not None: progress_total = target_total_steps else: ga = max(1, self.args.gradient_accumulation_steps) steps_per_epoch = max(len(self.train_dataloader) // ga, 1) progress_total = max(1, int(self.args.num_train_epochs)) * steps_per_epoch # Initialize Weights & Biases (rank0 only) if self._is_main(): if not hasattr(self, "_wandb"): proj = os.environ.get("WANDB_PROJECT", "codontranslator") name = os.environ.get("WANDB_NAME") run_id = os.environ.get("WANDB_RUN_ID") resume = os.environ.get("WANDB_RESUME") wandb_dir = os.environ.get("WANDB_DIR") world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else int(os.environ.get("WORLD_SIZE", "1")) init_kwargs = { "project": proj, "name": name, "config": { "lr": self.args.learning_rate, "warmup_ratio": self.args.warmup_ratio, "scheduler": self.args.lr_scheduler_type, "batch_size": self.args.per_device_train_batch_size, "eval_batch_size": self.args.per_device_eval_batch_size, "grad_accum": self.args.gradient_accumulation_steps, "effective_global_batch": self.args.per_device_train_batch_size * max(1, world_size) * max(1, self.args.gradient_accumulation_steps), "epochs": self.args.num_train_epochs, "steps_per_epoch": getattr(self.args, "steps_per_epoch", 0), "max_steps": self.args.max_steps, "weight_decay": self.args.weight_decay, "world_size": world_size, "output_dir": self.args.output_dir, "fsdp": self.args.fsdp, "bf16": self.args.bf16, "fp16": self.args.fp16, }, } if run_id: init_kwargs["id"] = run_id if resume: init_kwargs["resume"] = resume if wandb_dir: init_kwargs["dir"] = wandb_dir self._wandb = wandb.init(**init_kwargs) self.model.train() grad_accum = max(1, self.args.gradient_accumulation_steps) progress = None if self._is_main() and progress_total is not None and progress_total > 0: progress = tqdm(total=progress_total, initial=int(self.state["global_step"]), desc="Train", dynamic_ncols=True) if self.device.type == "cuda" and torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats(self.device) world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else int(os.environ.get("WORLD_SIZE", "1")) seqs_per_optimizer_step = ( int(self.args.per_device_train_batch_size) * max(1, world_size) * grad_accum ) log_window_start = time.perf_counter() log_window_optimizer_steps = 0 for epoch in range(self.state["epoch"], max(1, self.args.num_train_epochs)): self.state["epoch"] = epoch running_loss = 0.0 running_count = 0 train_iter = iter(self.train_dataloader) step = 0 batches_this_epoch = 0 optimizer_steps_this_epoch = 0 # If this is a streaming dataset with a shaped schedule, enforce a per-epoch optimizer step budget enforce_budget = False epoch_budget = None ds = self.train_dataloader.dataset if isinstance(ds, IterableDataset): spe = int(getattr(self.args, "steps_per_epoch", 0) or 0) if spe > 0: enforce_budget = True epoch_budget = int(spe) refill_attempts = 0 max_refills = 64 # avoids infinite loops when dataset is empty while True: batch, has_batch, local_has_batch = self._next_batch_sync(train_iter) if not has_batch: # If budget-enforced, attempt to refill the iterator and continue until budget is met. if enforce_budget and (epoch_budget is not None) and (optimizer_steps_this_epoch < epoch_budget): if local_has_batch and self._is_main(): logger.warning("Rank retained extra batch while peers exhausted stream; dropping to stay in sync") self._barrier() train_iter = iter(self.train_dataloader) refill_attempts += 1 if refill_attempts > max_refills: if self._is_main(): logger.warning( "Exceeded max refills for epoch %s (steps %s/%s). Ending epoch early.", epoch, optimizer_steps_this_epoch, epoch_budget, ) break continue else: if local_has_batch and self._is_main(): logger.warning("Rank retained extra batch while peers exhausted stream; dropping to stay in sync") break batch = self._prepare_batch(batch) batches_this_epoch += 1 codon_ids = batch["codon_ids"].to(self.device) input_ids = codon_ids[:, :-1] labels = codon_ids[:, :-1] # Mask PAD/EOS in labels pad_id = int(self.tokenizer.pad_token_id) if self.tokenizer is not None else 0 eos_id = int(self.tokenizer.special_ids.eos) if self.tokenizer is not None else -999 labels = labels.clone() labels[labels == pad_id] = -100 labels[labels == eos_id] = -100 cond = self._build_cond(batch) # autocast context use_cuda = (self.device.type == "cuda") autocast_dtype = self._amp_dtype if autocast_dtype is not None and use_cuda: ctx = torch.amp.autocast(device_type="cuda", dtype=autocast_dtype) else: from contextlib import nullcontext ctx = nullcontext() with ctx: out = self.model(codon_ids=input_ids, cond=cond, labels=labels, return_dict=True) loss = out["loss"] if self._scaler.is_enabled(): self._scaler.scale(loss / grad_accum).backward() else: (loss / grad_accum).backward() running_loss += float(loss.detach().item()) running_count += 1 do_step = ((step + 1) % grad_accum == 0) if do_step: # Clip if self.args.max_grad_norm and self.args.max_grad_norm > 0: if isinstance(self.model, FSDP): FSDP.clip_grad_norm_(self.model, self.args.max_grad_norm) else: if self._scaler.is_enabled(): self._scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.max_grad_norm) # Step if self._scaler.is_enabled(): self._scaler.step(self.optimizer) self._scaler.update() else: self.optimizer.step() if self.lr_scheduler is not None: self.lr_scheduler.step() self.optimizer.zero_grad(set_to_none=True) self.state["global_step"] += 1 optimizer_steps_this_epoch += 1 log_window_optimizer_steps += 1 # (wandb) Defer logging to the periodic block below # Log should_log = (self.state["global_step"] % max(1, self.args.logging_steps) == 0) peak_alloc_gb = 0.0 peak_reserved_gb = 0.0 if should_log: peak_alloc_gb, peak_reserved_gb = self._max_cuda_peak_gb() if self._is_main() and should_log: avg = running_loss / max(running_count, 1) lr = float(self.optimizer.param_groups[0]["lr"]) log_epoch = self._epoch_for_logging() elapsed = max(time.perf_counter() - log_window_start, 1e-9) step_time_s = elapsed / max(log_window_optimizer_steps, 1) seq_per_s = (seqs_per_optimizer_step * max(log_window_optimizer_steps, 1)) / elapsed msg = f"epoch {log_epoch} step {self.state['global_step']}: loss={avg:.4f} lr={lr:.6g}" if isinstance(out, dict): pl = out.get("prefix_len") pc = out.get("per_cap") if pl is not None and pc is not None: msg += f" prefix_mean={float(pl.detach().float().mean().item()):.1f} cap_mean={float(pc.detach().float().mean().item()):.1f}" msg += ( f" step_time_s={step_time_s:.3f} seq_per_s={seq_per_s:.1f}" f" peak_mem_alloc_gb={peak_alloc_gb:.1f} peak_mem_reserved_gb={peak_reserved_gb:.1f}" ) logger.info(msg) if hasattr(self, "_wandb"): wandb.log({ "train/loss": float(avg), "train/lr": float(lr), "perf/step_time_s": float(step_time_s), "perf/seq_per_s": float(seq_per_s), "system/peak_mem_alloc_gb": float(peak_alloc_gb), "system/peak_mem_reserved_gb": float(peak_reserved_gb), }, step=self.state["global_step"]) running_loss = 0.0 running_count = 0 log_window_start = time.perf_counter() log_window_optimizer_steps = 0 # Update progress bar if progress is not None: progress.update(1) # Stop when budget is reached for streaming schedule if target_total_steps is not None and self.state["global_step"] >= target_total_steps: metrics = {"train_loss": running_loss / max(running_count, 1)} self._save_checkpoint("final_model") self._barrier() return metrics # Periodic teacher-forced evaluation on the held-out dataset should_eval = ( self.eval_dataloader is not None and self.args.eval_interval > 0 and (self.state["global_step"] % self.args.eval_interval == 0) ) if should_eval: eval_metrics = self.evaluate() if self._is_main(): el = float(eval_metrics.get("eval_loss", 0.0)) ea = eval_metrics.get("eval_codon_acc", None) aa = eval_metrics.get("eval_aa_acc", None) if ea is not None and aa is not None: logger.info(f"eval: loss={el:.4f} codon_acc={float(ea):.3f} aa_acc={float(aa):.3f}") elif ea is not None: logger.info(f"eval: loss={el:.4f} codon_acc={float(ea):.3f}") elif aa is not None: logger.info(f"eval: loss={el:.4f} aa_acc={float(aa):.3f}") else: logger.info(f"eval: loss={el:.4f}") if hasattr(self, "_wandb"): log_payload = {"eval/loss": el} if ea is not None: log_payload["eval/codon_acc"] = float(ea) if aa is not None: log_payload["eval/aa_acc"] = float(aa) wandb.log(log_payload, step=self.state["global_step"]) # Save by step if self.args.save_steps > 0 and (self.state["global_step"] % self.args.save_steps == 0): self._save_checkpoint(f"checkpoint-{self.state['global_step']}") # Hard horizon for streaming/step-limited runs if self.args.max_steps > 0 and self.state["global_step"] >= self.args.max_steps: metrics = {"train_loss": running_loss / max(running_count, 1)} self._save_checkpoint("final_model") self._barrier() if progress is not None: progress.close() return metrics step += 1 # If we enforce a per-epoch budget for streaming datasets, end the epoch once it's reached if enforce_budget and (epoch_budget is not None) and (optimizer_steps_this_epoch >= epoch_budget): break # Epoch summary (rank0 only) if self._is_main(): try: eb = int(epoch_budget) if epoch_budget is not None else -1 except Exception: eb = -1 logger.info( "epoch %s completed: optimizer_steps=%s%s", self._epoch_for_logging(), optimizer_steps_this_epoch, (f" / budget {eb}" if eb > 0 else ""), ) if dist.is_available() and dist.is_initialized(): gather_device = self.device if self.device.type == "cuda" else torch.device("cpu") counts_tensor = torch.tensor( [batches_this_epoch, optimizer_steps_this_epoch], dtype=torch.long, device=gather_device, ) gathered = [torch.zeros_like(counts_tensor) for _ in range(dist.get_world_size())] dist.all_gather(gathered, counts_tensor) batch_counts = [int(t[0].item()) for t in gathered] step_counts = [int(t[1].item()) for t in gathered] batch_gap = max(batch_counts) - min(batch_counts) step_gap = max(step_counts) - min(step_counts) if self._is_main() and (batch_gap > 0 or step_gap > 0): logger.warning( "Epoch %s imbalance detected across ranks: batches min=%s max=%s, optimizer steps min=%s max=%s", epoch, min(batch_counts), max(batch_counts), min(step_counts), max(step_counts), ) # Epoch boundary save for sized datasets if not isinstance(ds, IterableDataset): self._save_checkpoint(f"epoch-{epoch}") metrics = {"train_loss": 0.0} if progress is not None: progress.close() self._barrier() return metrics # ---- evaluation ---- def evaluate(self) -> Dict[str, float]: if self.eval_dataloader is None: return {"eval_loss": 0.0} self.model.eval() loss_sum = 0.0 loss_tokens = 0 codon_correct = 0 codon_total = 0 aa_correct = 0 aa_total = 0 tok = self.tokenizer pad_id = int(tok.pad_token_id) if tok is not None else 0 eos_id = int(tok.special_ids.eos) if tok is not None and hasattr(tok, "special_ids") else -999 num_special = int(tok.num_special_tokens) if tok is not None else 0 codon2aa = tok.codon2aa_char_map() if tok is not None and hasattr(tok, "codon2aa_char_map") else {} is_streaming = isinstance(self.eval_dataloader.dataset, IterableDataset) max_batches = int(self.args.eval_steps) if (is_streaming and self.args.eval_steps > 0) else None with torch.no_grad(): eval_iter = iter(self.eval_dataloader) b_idx = 0 while True: batch, has_batch, local_has_batch = self._next_batch_sync(eval_iter) if not has_batch: if local_has_batch and self._is_main(): logger.debug("eval dataloader: discarded tail batch to stay in sync across ranks") break if max_batches is not None and b_idx >= max_batches: break batch = self._prepare_batch(batch) codon_ids = batch["codon_ids"].to(self.device) input_ids = codon_ids[:, :-1] labels = codon_ids[:, :-1] labels = labels.clone() labels[labels == pad_id] = -100 labels[labels == eos_id] = -100 cond = self._build_cond(batch) use_cuda = (self.device.type == "cuda") autocast_dtype = self._amp_dtype if autocast_dtype is not None and use_cuda: ctx = torch.amp.autocast(device_type="cuda", dtype=autocast_dtype) else: from contextlib import nullcontext ctx = nullcontext() with ctx: out = self.model(codon_ids=input_ids, cond=cond, labels=labels, return_dict=True) loss = out.get("loss") per_cap = out.get("per_cap") logits = out.get("logits") tokens_in_batch = 0 if per_cap is not None: tokens_in_batch = int(torch.clamp(per_cap.detach(), min=0).sum().item()) loss_tokens += tokens_in_batch if loss is not None and tokens_in_batch > 0: loss_sum += float(loss.detach().item()) * tokens_in_batch if logits is None or logits.size(1) == 0 or per_cap is None: continue max_cap = logits.size(1) batch_size = logits.size(0) labels_aligned = torch.full((batch_size, max_cap), -100, dtype=labels.dtype, device=labels.device) common_cols = min(labels.size(1), max_cap) if common_cols > 0: labels_aligned[:, :common_cols] = labels[:, :common_cols] per_cap_int = torch.clamp(per_cap.to(dtype=torch.long), min=0, max=max_cap) for row in range(batch_size): cap = int(per_cap_int[row].item()) if cap < max_cap: labels_aligned[row, cap:] = -100 supervised = labels_aligned != -100 if num_special > 0: supervised = supervised & (labels_aligned >= num_special) if not supervised.any(): continue preds = logits.argmax(dim=-1) codon_correct += int((preds[supervised] == labels_aligned[supervised]).sum().item()) codon_total += int(supervised.sum().item()) if codon2aa and isinstance(batch, dict) and "protein_seqs" in batch: prot_list = batch.get("protein_seqs", []) for row in range(batch_size): cap = int(per_cap_int[row].item()) if cap <= 0: continue mask_row = supervised[row, :cap] if not mask_row.any(): continue preds_row = preds[row, :cap][mask_row] prot = prot_list[row] if (isinstance(prot_list, list) and row < len(prot_list)) else "" if not prot: continue seq_len = min(len(prot), preds_row.size(0)) if seq_len <= 0: continue pred_aa = ''.join(codon2aa.get(int(t.item()), 'X') for t in preds_row[:seq_len]) truth_aa = prot[:seq_len] aa_correct += sum(1 for i in range(seq_len) if pred_aa[i] == truth_aa[i]) aa_total += seq_len b_idx += 1 totals = torch.tensor( [loss_sum, loss_tokens, codon_correct, codon_total, aa_correct, aa_total], dtype=torch.float64, device=self.device, ) if dist.is_available() and dist.is_initialized(): # Ensure every rank has finished its forward passes before the final # metric reduction, otherwise FSDP may still be issuing _all_gather # collectives on slower ranks. self._barrier() dist.all_reduce(totals, op=dist.ReduceOp.SUM) loss_sum, loss_tokens, codon_correct, codon_total, aa_correct, aa_total = totals.tolist() self.model.train() metrics: Dict[str, float] = {"eval_loss": float(loss_sum) / loss_tokens if loss_tokens > 0 else 0.0} if codon_total > 0: metrics["eval_codon_acc"] = float(codon_correct) / codon_total if aa_total > 0: metrics["eval_aa_acc"] = float(aa_correct) / aa_total self._barrier() return metrics # ---- internals ---- def _setup_fsdp(self): # Ensure default process group is initialized (required by FSDP) device = self.device if dist.is_available() and not dist.is_initialized(): backend = "nccl" if device.type == "cuda" else "gloo" sig = inspect.signature(dist.init_process_group) if "timeout" in sig.parameters: dist.init_process_group(backend=backend, init_method="env://", timeout=datetime.timedelta(minutes=30)) else: dist.init_process_group(backend=backend, init_method="env://") mp = MixedPrecision( param_dtype=(torch.float16 if self.args.fp16 else torch.bfloat16 if self.args.bf16 else torch.float32), reduce_dtype=(torch.float16 if self.args.fp16 else torch.bfloat16 if self.args.bf16 else torch.float32), buffer_dtype=torch.float32, ) logger.info(f"FSDP enabled: sharding={self.args.fsdp} mp_param={mp.param_dtype} mp_reduce={mp.reduce_dtype}") # Keep frozen ESM off FSDP if present base = self._unwrap(self.model) ignored = [] if hasattr(base, "esm") and isinstance(base.esm, nn.Module): ignored.append(base.esm) self.model = FSDP( self.model, device_id=(self.device if device.type == "cuda" else None), sharding_strategy=ShardingStrategy.FULL_SHARD, mixed_precision=mp, ignored_modules=(ignored if ignored else None), sync_module_states=True, ) # Place ignored module on device exactly once if ignored: ignored[0].to(device) def _unwrap(self, module): return getattr(module, "module", module) def _prepare_batch(self, batch: Dict[str, Any]) -> Dict[str, Any]: # Species embeddings (fixed-size or sequence) if self.species_store is not None and "species_ids" in batch: sids = batch["species_ids"] if torch.is_tensor(sids): sids = sids.detach().cpu().tolist() result = self.species_store.batch_get(sids) if isinstance(result, tuple): sp_tok, _ = result # [B, Ls, Ds] batch["species_tok_emb"] = sp_tok.to(self.device, non_blocking=True) else: sp = result # [B, Ds] batch["species_emb"] = sp.to(self.device, non_blocking=True) # Move obvious tensors if "codon_ids" in batch and hasattr(batch["codon_ids"], "to"): batch["codon_ids"] = batch["codon_ids"].to(self.device, non_blocking=True) return batch def _build_cond(self, batch: Dict[str, Any]) -> Dict[str, Any]: cond: Dict[str, Any] = {"control_mode": "fixed"} if "species_tok_emb" in batch: cond["species_tok_emb_src"] = batch["species_tok_emb"] cond["species_tok_emb_tgt"] = batch["species_tok_emb"] elif "species_emb" in batch: cond["species_emb_src"] = batch["species_emb"] cond["species_emb_tgt"] = batch["species_emb"] if "protein_seqs" in batch: cond["protein_seqs"] = batch["protein_seqs"] return cond def _next_batch_sync(self, iterator): """Fetch next batch and drop out early if any rank exhausts its loader.""" try: batch = next(iterator) local_has_batch = True except StopIteration: batch = None local_has_batch = False distributed = dist.is_available() and dist.is_initialized() has_batch = local_has_batch if distributed: flag_device = self.device if self.device.type == "cuda" else torch.device("cpu") flag = torch.tensor([1 if local_has_batch else 0], device=flag_device) dist.all_reduce(flag, op=dist.ReduceOp.MIN) has_batch = bool(flag.item()) if not has_batch: return None, False, local_has_batch return batch, True, local_has_batch def _is_main(self) -> bool: return (not dist.is_available()) or (not dist.is_initialized()) or dist.get_rank() == 0 def _barrier(self): if dist.is_available() and dist.is_initialized(): # On NCCL, pass device_ids to avoid rank↔GPU mapping ambiguity when supported if self.device.type == "cuda": sig = inspect.signature(dist.barrier) if "device_ids" in sig.parameters: dist.barrier(device_ids=[self.local_rank]) return dist.barrier() def _max_cuda_peak_gb(self) -> Tuple[float, float]: if self.device.type != "cuda" or not torch.cuda.is_available(): return 0.0, 0.0 vals = torch.tensor( [ float(torch.cuda.max_memory_allocated(self.device)), float(torch.cuda.max_memory_reserved(self.device)), ], dtype=torch.float64, device=self.device, ) if dist.is_available() and dist.is_initialized(): dist.all_reduce(vals, op=dist.ReduceOp.MAX) scale = float(1024 ** 3) return float(vals[0].item() / scale), float(vals[1].item() / scale) # (Per-sample quick eval removed; evaluation now uses held-out dataloader.) def _epoch_for_logging(self) -> int: steps_per_epoch = int(getattr(self.args, "steps_per_epoch", 0) or 0) if steps_per_epoch > 0: est = self.state.get("global_step", 0) // steps_per_epoch if self.args.num_train_epochs > 0: max_epoch = max(int(self.args.num_train_epochs) - 1, 0) if est > max_epoch: return max_epoch return int(est) return int(self.state.get("epoch", 0)) # ---- checkpointing ---- def _save_checkpoint(self, name: str): self.state["epoch"] = int(self._epoch_for_logging()) # All ranks participate in FSDP state_dict collectives; only rank0 writes files. out_dir = os.path.join(self.args.output_dir, name) os.makedirs(out_dir, exist_ok=True) optim_state = None if isinstance(self.model, FSDP): with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=FutureWarning) with FSDP.state_dict_type( self.model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(rank0_only=True, offload_to_cpu=True), FullOptimStateDictConfig(rank0_only=True, offload_to_cpu=True), ): state = self.model.state_dict() # NOTE: Under FSDP, optimizer.state_dict() is sharded per-rank. # Use FSDP.optim_state_dict() to materialize a full optimizer state dict (rank0_only). if self.optimizer is not None: optim_state = FSDP.optim_state_dict(self.model, self.optimizer) else: state = self._unwrap(self.model).state_dict() if self.optimizer is not None: optim_state = self.optimizer.state_dict() # Save minimal data cursor (total samples yielded so far) next to output_dir if configured per_rank_positions: Optional[List[int]] = None p = getattr(self.args, "data_cursor_path", None) if p: ds = getattr(self.train_dataloader, "dataset", None) if hasattr(ds, "get_stream_position"): local_pos = int(ds.get_stream_position()) if dist.is_available() and dist.is_initialized(): gather_device = self.device if self.device.type == "cuda" else torch.device("cpu") tensor = torch.tensor([local_pos], dtype=torch.long, device=gather_device) gathered = [torch.zeros_like(tensor) for _ in range(dist.get_world_size())] dist.all_gather(gathered, tensor) per_rank_positions = [int(t.item()) for t in gathered] else: per_rank_positions = [local_pos] if not self._is_main(): # Non-main ranks skip serialization but stay in lockstep self._barrier() return # Rank 0 writes artifacts save_file(state, os.path.join(out_dir, "model.safetensors")) # Optimizer + scheduler if optim_state is not None: torch.save(optim_state, os.path.join(out_dir, "optimizer.pt")) if self.lr_scheduler is not None: torch.save(self.lr_scheduler.state_dict(), os.path.join(out_dir, "scheduler.pt")) # Trainer config/state base = self._unwrap(self.model) # Infer mlp_ratio from first block if present mlp_ratio = 4.0 try: if hasattr(base, "blocks") and len(getattr(base, "blocks", [])) > 0: w1 = base.blocks[0].ffn.w1.weight # [H*mlp, H] H = int(getattr(base, "hidden_size", w1.shape[1])) if H > 0: mlp_ratio = float(w1.shape[0]) / float(H) except Exception: pass trainer_cfg = { # capacity / prefixes "max_length": int(self.args.max_length), "max_species_prefix": int(getattr(base, "max_species_prefix", 0)), "max_protein_prefix": int(getattr(base, "max_protein_prefix", 0)), # architecture hints "hidden_size": int(getattr(base, "hidden_size", 0)), "num_hidden_layers": int(getattr(base, "num_layers", 0)), "num_attention_heads": int(getattr(base, "num_heads", 0)), "mlp_ratio": float(mlp_ratio), # conditioning flags "prepend_species": bool(getattr(base, "prepend_species", True)), "prepend_protein": bool(getattr(base, "prepend_protein", False)), "species_embedding_dim": int(getattr(base, "species_embedding_dim", 1024)), # ESM info (even if prepend_protein=False) "esm_model_name": str(getattr(self.args, "esm_model_name", "")), "esm_device": str(getattr(self.args, "esm_device", "cuda")), "esm_dtype": str(getattr(self.args, "esm_dtype", "fp32")).lower(), # kernels # attention impl "attn_impl": str(getattr(base, "attn_impl", "gqa")), "num_kv_groups": int(getattr(base, "num_kv_groups", 0)), } with open(os.path.join(out_dir, "trainer_config.json"), "w") as f: json.dump(trainer_cfg, f, indent=2) with open(os.path.join(out_dir, "trainer_state.json"), "w") as f: json.dump({"epoch": self.state["epoch"], "global_step": self.state["global_step"]}, f, indent=2) if p and per_rank_positions is not None: payload = { "skip_samples": int(sum(per_rank_positions)), "per_rank": per_rank_positions, "world_size": len(per_rank_positions), } os.makedirs(os.path.dirname(os.path.abspath(p)), exist_ok=True) with open(p, "w") as f: json.dump(payload, f) # Tokenizer vocab for sampling try: if self.tokenizer is not None and hasattr(self.tokenizer, "save_vocabulary"): self.tokenizer.save_vocabulary(out_dir) except Exception as e: logger.warning(f"Failed to save vocabulary to {out_dir}: {e}") self._prune_checkpoints(self.args.output_dir, self.args.save_total_limit) logger.info(f"Saved checkpoint → {out_dir}") # Release other ranks self._barrier() def _resume_from(self, ckpt_dir: str): st_path = os.path.join(ckpt_dir, "model.safetensors") if not os.path.exists(st_path): raise FileNotFoundError(f"No model.safetensors in {ckpt_dir}") state = load_file(st_path) if isinstance(self.model, FSDP): with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=FutureWarning) with FSDP.state_dict_type( self.model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(rank0_only=False, offload_to_cpu=True), ): self.model.load_state_dict(state, strict=False) else: self._unwrap(self.model).load_state_dict(state, strict=False) scheduler_restored = False opt_path = os.path.join(ckpt_dir, "optimizer.pt") if os.path.exists(opt_path): if self.optimizer is None: self._create_optimizer_and_scheduler() if not self.args.override_lr_on_resume: loaded = torch.load(opt_path, map_location="cpu") # Under FSDP, saved optimizer.pt is a full optimizer state dict produced by # FSDP.optim_state_dict(). Convert it to a per-rank state dict before loading. if isinstance(self.model, FSDP): try: loaded = FSDP.optim_state_dict_to_load(self.model, self.optimizer, loaded) except Exception as e: msg = ( "Failed to convert FSDP optimizer state dict for loading. " "This checkpoint likely contains an incomplete (rank0-only sharded) optimizer.pt from an older version. " "Full optimizer resume is not possible from this checkpoint.\n" f"Underlying error: {e}\n" "Options:\n" " 1) Start a fresh run (new --output_dir), or\n" " 2) Re-run with --override_lr_on_resume to skip optimizer restore (not a full resume)." ) if self._is_main(): logger.error(msg) raise RuntimeError(msg) from e self.optimizer.load_state_dict(loaded) sch_path = os.path.join(ckpt_dir, "scheduler.pt") if os.path.exists(sch_path): if self.lr_scheduler is None: self._create_optimizer_and_scheduler() if self.lr_scheduler is not None and not self.args.override_lr_on_resume: self.lr_scheduler.load_state_dict(torch.load(sch_path, map_location="cpu")) scheduler_restored = True ts_path = os.path.join(ckpt_dir, "trainer_state.json") if os.path.exists(ts_path): with open(ts_path, "r") as f: ts = json.load(f) self.state["epoch"] = int(ts.get("epoch", 0)) self.state["global_step"] = int(ts.get("global_step", 0)) steps_per_epoch = int(getattr(self.args, "steps_per_epoch", 0) or 0) if steps_per_epoch > 0: inferred_epoch = self.state.get("global_step", 0) // steps_per_epoch num_epochs = max(int(self.args.num_train_epochs), 1) inferred_epoch = min(inferred_epoch, num_epochs - 1) if inferred_epoch != self.state.get("epoch"): if self._is_main(): logger.info( "Adjusting epoch from %s to %s based on global_step %s and steps_per_epoch %s", self.state.get("epoch"), inferred_epoch, self.state.get("global_step"), steps_per_epoch, ) self.state["epoch"] = int(inferred_epoch) # If we skipped loading the scheduler state (e.g., different world size or override), # fast-forward it to the saved global_step so LR does not restart from warmup. if self.lr_scheduler is not None and not scheduler_restored: target_step = int(self.state.get("global_step", 0)) if target_step > 0: try: # Most schedulers (LambdaLR, CosineAnnealing, etc.) accept an "epoch" kwarg. self.lr_scheduler.step(target_step) except TypeError: # Fallback: advance manually. for _ in range(target_step): self.lr_scheduler.step() # Ensure optimizer LR reflects the scheduler's current value. try: last_lrs = self.lr_scheduler.get_last_lr() except Exception: last_lrs = [group.get("lr") for group in self.optimizer.param_groups] if last_lrs: for group, lr in zip(self.optimizer.param_groups, last_lrs): group["lr"] = float(lr) logger.info(f"Resumed from {ckpt_dir}") def _checkpoint_step(self, path: str) -> Optional[int]: m = re.fullmatch(r"checkpoint-(\d+)", os.path.basename(path)) if not m: return None return int(m.group(1)) def _prune_checkpoints(self, root: str, keep: int): if not os.path.isdir(root): return try: subdirs = [ os.path.join(root, d) for d in os.listdir(root) if os.path.isdir(os.path.join(root, d)) ] except FileNotFoundError: return step_dirs: list[tuple[int, str]] = [] for path in subdirs: step = self._checkpoint_step(path) if step is not None: step_dirs.append((step, path)) if not step_dirs: return step_dirs.sort(key=lambda item: item[0]) latest_step = step_dirs[-1][0] recent_window = max(0, int(getattr(self.args, "ckpt_recent_window_steps", 0) or 0)) recent_interval = max(0, int(getattr(self.args, "ckpt_recent_interval", 0) or 0)) archive_interval = max(0, int(getattr(self.args, "ckpt_archive_interval", 0) or 0)) keep_paths: set[str] = set() if recent_window > 0 and (recent_interval > 0 or archive_interval > 0): if recent_interval <= 0: recent_interval = max(1, int(getattr(self.args, "save_steps", 1) or 1)) for step, path in step_dirs: age = latest_step - step if age <= recent_window: interval = recent_interval else: interval = archive_interval if interval > 0 and (step % interval == 0): keep_paths.add(path) if not keep_paths: # Legacy fallback: keep the most recent N step checkpoints. if keep <= 0: return keep_paths = {path for _, path in step_dirs[-keep:]} else: # Always preserve the newest checkpoint, even if the interval math misses it. keep_paths.add(step_dirs[-1][1]) if keep > 0: kept = [(step, path) for step, path in step_dirs if path in keep_paths] if len(kept) > keep: trim = len(kept) - keep for _, path in kept[:trim]: keep_paths.discard(path) removed = [] for _, path in step_dirs: if path in keep_paths: continue shutil.rmtree(path, ignore_errors=True) removed.append(os.path.basename(path)) if removed and self._is_main(): logger.info( "Pruned %s checkpoints (latest_step=%s, recent_window=%s, recent_interval=%s, archive_interval=%s)", len(removed), latest_step, recent_window, recent_interval, archive_interval, )