from __future__ import annotations import argparse import os import random import sys from dataclasses import dataclass from pathlib import Path from typing import Any PROJECT_ROOT = Path(__file__).resolve().parent.parent MODEL_ROOT = PROJECT_ROOT / "model" SCRIPT_ROOT = PROJECT_ROOT / "scripts" for path in (PROJECT_ROOT, MODEL_ROOT, SCRIPT_ROOT): if str(path) not in sys.path: sys.path.insert(0, str(path)) import torch import torch.distributed as dist import torch.nn.functional as F import yaml from torch import nn from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler from stormer import build_diffusion_model, build_regression_model from data_loader import StormCastDataset @dataclass class DistributedContext: device: torch.device rank: int local_rank: int world_size: int @property def distributed(self) -> bool: return self.world_size > 1 @property def is_main(self) -> bool: return self.rank == 0 def regression_loss( model: nn.Module, condition: torch.Tensor, target: torch.Tensor, ) -> torch.Tensor: return F.mse_loss(model(condition), target) def edm_residual_loss( model: nn.Module, residual: torch.Tensor, condition: torch.Tensor, sigma_data: float = 0.5, p_mean: float = -1.2, p_std: float = 1.2, ) -> torch.Tensor: sigma = torch.exp( torch.randn(residual.shape[0], device=residual.device) * p_std + p_mean ) noise = torch.randn_like(residual) * sigma[:, None, None, None] denoised = model(residual + noise, sigma, condition=condition) weight = (sigma.square() + sigma_data**2) / (sigma * sigma_data).square() return (weight[:, None, None, None] * (denoised - residual).square()).mean() def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Train the two-stage StormCast model") parser.add_argument("--config", type=Path, default=Path("conf/config.yaml")) parser.add_argument("--stage", choices=("regression", "diffusion")) parser.add_argument("--resume", type=Path) parser.add_argument("--initial-weights", type=Path) parser.add_argument("--regression-weights", type=Path) parser.add_argument("--max-steps", type=int) return parser.parse_args() def main() -> None: args = parse_args() config_path = args.config.resolve() with config_path.open("r", encoding="utf-8") as handle: config = yaml.safe_load(handle) _resolve_config_paths(config, config_path.parent.parent) training = config["training"] stage = args.stage or training["stage"] resume = args.resume or _optional_path(training.get("resume_checkpoint")) initial_weights = args.initial_weights or _optional_path(training.get("initial_weights")) regression_weights = args.regression_weights or _optional_path( training.get("regression_weights") ) max_steps = args.max_steps if args.max_steps is not None else training["max_steps"] if resume is None and initial_weights is None and not training["from_scratch"]: checkpoint_key = f"{stage}_checkpoint" initial_weights = _optional_path(config["model"].get(checkpoint_key)) if initial_weights is None: raise ValueError( f"training.from_scratch is false but model.{checkpoint_key} is not set" ) context = initialize_distributed() _seed_everything(config["project"]["seed"], context.rank) try: train( config=config, stage=stage, context=context, resume=resume, initial_weights=initial_weights, regression_weights=regression_weights, max_steps=max_steps, ) finally: if dist.is_initialized(): dist.destroy_process_group() def train( config: dict[str, Any], stage: str, context: DistributedContext, resume: Path | None, initial_weights: Path | None, regression_weights: Path | None, max_steps: int | None, ) -> None: if stage not in ("regression", "diffusion"): raise ValueError("training.stage must be 'regression' or 'diffusion'") data_config = config["data"] loader_config = config["dataloader"] training_config = config["training"] if list(data_config["image_size"]) != list(config["model"]["image_size"]): raise ValueError("Data and model image sizes must match") if list(data_config["era5_image_size"]) != [721, 1440]: raise ValueError("ERA5 grid must be 721 x 1440") dataset = StormCastDataset( data_root=data_config["root_dir"], years=data_config["train_years"], era5_variables=data_config["era5_variables"], state_variables=data_config["state_variables"], invariant_variables=data_config["invariant_variables"], image_size=data_config["image_size"], input_steps=data_config["input_steps"], output_steps=data_config["output_steps"], normalize=data_config["normalize"], ) sampler = ( DistributedSampler( dataset, num_replicas=context.world_size, rank=context.rank, shuffle=True, ) if context.distributed else None ) loader = DataLoader( dataset, batch_size=loader_config["batch_size"], shuffle=sampler is None, sampler=sampler, num_workers=loader_config["num_workers"], pin_memory=loader_config["pin_memory"], drop_last=False, ) regression, model = _build_stage_models( config, stage, regression_weights, context.device ) if initial_weights is not None and resume is None: _load_initial_weights(model, initial_weights, stage) model.to(context.device) optimizer = torch.optim.Adam( model.parameters(), lr=training_config["learning_rate"], betas=tuple(training_config["betas"]), weight_decay=training_config["weight_decay"], ) start_epoch = 0 start_batch = 0 global_step = 0 if resume is not None: start_epoch, start_batch, global_step = load_training_checkpoint( resume, model, optimizer, stage, context.device ) if context.distributed: model = DistributedDataParallel( model, device_ids=[context.local_rank], output_device=context.local_rank, ) checkpoint_dir = Path(training_config["checkpoint_dir"]) / stage if context.is_main: checkpoint_dir.mkdir(parents=True, exist_ok=True) print( f"stage={stage} world_size={context.world_size} " f"parameters={sum(parameter.numel() for parameter in model.parameters())}" ) stop = False if max_steps is not None and global_step >= max_steps: stop = True for epoch in range(start_epoch, training_config["epochs"]): if stop: break if sampler is not None: sampler.set_epoch(epoch) model.train() for batch_index, batch in enumerate(loader): if epoch == start_epoch and batch_index < start_batch: continue loss = _training_step( stage, model, regression, batch, context.device, training_config, ) optimizer.zero_grad(set_to_none=True) loss.backward() clip_norm = training_config.get("gradient_clip_norm") if clip_norm is not None: torch.nn.utils.clip_grad_norm_(model.parameters(), clip_norm) optimizer.step() global_step += 1 reduced_loss = _mean_across_ranks(loss.detach(), context.world_size) if context.is_main and global_step % training_config["log_interval"] == 0: print( f"stage={stage} epoch={epoch + 1} step={global_step} " f"loss={reduced_loss.item():.8f}" ) if ( context.is_main and global_step % training_config["checkpoint_interval"] == 0 ): save_training_checkpoint( checkpoint_dir / "model_bak.pt", model, optimizer, stage, epoch, batch_index + 1, global_step, config, ) if max_steps is not None and global_step >= max_steps: stop = True break if stop: break if context.is_main: checkpoint = checkpoint_dir / "model_bak.pt" save_training_checkpoint( checkpoint, model, optimizer, stage, epoch if "epoch" in locals() else start_epoch, batch_index + 1 if "batch_index" in locals() else start_batch, global_step, config, ) print(f"checkpoint={checkpoint} steps={global_step}") def _build_stage_models( config: dict[str, Any], stage: str, regression_weights: Path | None, device: torch.device, ) -> tuple[nn.Module | None, nn.Module]: data_config = config["data"] model_config = config["model"] common = { "image_size": model_config["image_size"], "state_channels": len(data_config["state_variables"]), "invariant_channels": len(data_config["invariant_variables"]), "model_channels": model_config["model_channels"], "channel_mult": model_config["channel_mult"], "num_blocks": model_config["num_blocks"], "attn_resolutions": model_config["attention_resolutions"], } if stage == "regression": model = build_regression_model( **common, background_channels=len(data_config["era5_variables"]), ) return None, model if regression_weights is None: raise ValueError("Diffusion training requires --regression-weights") regression = _load_model_weights( build_regression_model( **common, background_channels=len(data_config["era5_variables"]), ), regression_weights, "regression", ).to(device) regression.eval() regression.requires_grad_(False) return regression, build_diffusion_model(**common) def _training_step( stage: str, model: nn.Module, regression: nn.Module | None, batch: dict[str, Any], device: torch.device, training_config: dict[str, Any], ) -> torch.Tensor: background = batch["background"].to(device, dtype=torch.float32) state, target = ( tensor.to(device, dtype=torch.float32) for tensor in batch["state"] ) invariant = batch["invariant"].to(device, dtype=torch.float32) if invariant.ndim == 3: invariant = invariant.unsqueeze(0) if invariant.shape[0] == 1 and state.shape[0] > 1: invariant = invariant.expand(state.shape[0], -1, -1, -1) if stage == "regression": condition = torch.cat((state, background, invariant), dim=1) return regression_loss(model, condition, target) if regression is None: raise RuntimeError("Regression model is required for diffusion training") with torch.no_grad(): regression_condition = torch.cat((state, background, invariant), dim=1) regression_prediction = regression(regression_condition) residual = target - regression_prediction condition = torch.cat((state, regression_prediction, invariant), dim=1) return edm_residual_loss( model, residual, condition, sigma_data=training_config["sigma_data"], p_mean=training_config["P_mean"], p_std=training_config["P_std"], ) def save_training_checkpoint( path: Path, model: nn.Module, optimizer: torch.optim.Optimizer, stage: str, epoch: int, batch_in_epoch: int, global_step: int, config: dict[str, Any], ) -> None: model = model.module if isinstance(model, DistributedDataParallel) else model temporary = path.with_suffix(path.suffix + ".tmp") torch.save( { "stage": stage, "epoch": epoch, "batch_in_epoch": batch_in_epoch, "global_step": global_step, "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "config": config, }, temporary, ) temporary.replace(path) def load_training_checkpoint( path: Path, model: nn.Module, optimizer: torch.optim.Optimizer, stage: str, device: torch.device, ) -> tuple[int, int, int]: checkpoint = torch.load(path, map_location=device, weights_only=False) if checkpoint["stage"] != stage: raise ValueError( f"Checkpoint stage is {checkpoint['stage']}, requested stage is {stage}" ) model.load_state_dict(checkpoint["model_state_dict"], strict=True) optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) return ( int(checkpoint["epoch"]), int(checkpoint.get("batch_in_epoch", 0)), int(checkpoint["global_step"]), ) def initialize_distributed() -> DistributedContext: world_size = int(os.environ.get("WORLD_SIZE", "1")) local_rank = int(os.environ.get("LOCAL_RANK", "0")) if not torch.cuda.is_available(): raise RuntimeError("StormCast training requires a CUDA/HIP device") torch.cuda.set_device(local_rank) if world_size > 1: dist.init_process_group(backend="nccl", init_method="env://") rank = dist.get_rank() else: rank = 0 return DistributedContext( device=torch.device("cuda", local_rank), rank=rank, local_rank=local_rank, world_size=world_size, ) def _load_initial_weights(model: nn.Module, path: Path, stage: str) -> None: loaded = _load_model_weights(model, path, stage) model.load_state_dict(loaded.state_dict(), strict=True) def _load_model_weights(model: nn.Module, path: Path, stage: str) -> nn.Module: checkpoint = torch.load(path, map_location="cpu", weights_only=False) state = checkpoint.get("model_state_dict", checkpoint) model.load_state_dict(state, strict=True) return model def _mean_across_ranks(value: torch.Tensor, world_size: int) -> torch.Tensor: if world_size > 1: dist.all_reduce(value, op=dist.ReduceOp.SUM) value /= world_size return value def _resolve_config_paths(config: dict[str, Any], project_root: Path) -> None: for section, key in ( ("data", "root_dir"), ("training", "checkpoint_dir"), ): path = Path(config[section][key]) if not path.is_absolute(): config[section][key] = str((project_root / path).resolve()) for key in ("regression_weights", "diffusion_weights"): value = config["model"].get(key) if value: path = Path(value) if not path.is_absolute(): config["model"][key] = str((project_root / path).resolve()) for key in ("initial_weights", "resume_checkpoint", "regression_weights"): value = config["training"].get(key) if value: path = Path(value) if not path.is_absolute(): config["training"][key] = str((project_root / path).resolve()) def _optional_path(value: str | Path | None) -> Path | None: return None if value is None else Path(value) def _seed_everything(seed: int, rank: int) -> None: seed += rank random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) if __name__ == "__main__": main()