| from __future__ import annotations |
|
|
| import argparse |
| import os |
| from pathlib import Path |
| import sys |
| import traceback |
| from typing import Any |
|
|
| import torch |
| import torch.distributed as dist |
| from torch.nn.parallel import DistributedDataParallel |
| from torch.utils.data import DataLoader |
| from torch.utils.data.distributed import DistributedSampler |
| import yaml |
|
|
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| if str(PROJECT_ROOT) not in sys.path: |
| sys.path.insert(0, str(PROJECT_ROOT)) |
| SCRIPTS_DIR = PROJECT_ROOT / "scripts" |
| if str(SCRIPTS_DIR) not in sys.path: |
| sys.path.insert(0, str(SCRIPTS_DIR)) |
|
|
| from model.checkpoint import load_checkpoint |
| from model.w_mae import w_mae_base |
| from era5_adapter import WMAEERA5Dataset |
|
|
|
|
| def load_config(path: Path) -> dict[str, Any]: |
| with path.open("r", encoding="utf-8") as handle: |
| config = yaml.safe_load(handle) |
| if not isinstance(config, dict): |
| raise ValueError(f"Configuration must contain a mapping: {path}") |
| return config |
|
|
|
|
| def resolve_project_path(value: str | Path, config_path: Path) -> Path: |
| path = Path(value) |
| return path if path.is_absolute() else (config_path.parent.parent / path).resolve() |
|
|
|
|
| def validate_config(config: dict[str, Any]) -> None: |
| data = config["data"] |
| model = config["model"] |
| if int(data["expected_channels"]) != int(model["input_channels"]): |
| raise ValueError("data.expected_channels must equal model.input_channels.") |
| if tuple(data["model_size"]) != tuple(model["image_size"]): |
| raise ValueError("data.model_size must equal model.image_size.") |
| if tuple(model["patch_size"])[0] <= 0 or tuple(model["patch_size"])[1] <= 0: |
| raise ValueError("model.patch_size values must be positive.") |
| if float(model["mask_ratio"]) not in {0.0, 0.75}: |
| raise ValueError("W-MAE supports mask_ratio 0.0 or 0.75 only.") |
|
|
|
|
| def build_model(config: dict[str, Any]) -> torch.nn.Module: |
| model_config = config["model"] |
| return w_mae_base( |
| img_size=tuple(model_config["image_size"]), |
| patch_size=tuple(model_config["patch_size"]), |
| in_chans=int(model_config["input_channels"]), |
| embed_dim=int(model_config["embed_dim"]), |
| depth=int(model_config["encoder_depth"]), |
| decoder_embed_dim=int(model_config["decoder_embed_dim"]), |
| decoder_depth=int(model_config["decoder_depth"]), |
| mlp_ratio=float(model_config.get("mlp_ratio", 4.0)), |
| norm_pix_loss=bool(model_config["norm_pixel_loss"]), |
| ) |
|
|
|
|
| def build_dataset(config: dict[str, Any], config_path: Path, split: str) -> WMAEERA5Dataset: |
| data = config["data"] |
| years = data[f"{split}_years"] |
| if not data["variables"]: |
| raise ValueError( |
| "data.variables is empty. Supply the verified, explicitly ordered 20-channel list before training." |
| ) |
| return WMAEERA5Dataset( |
| dataset_dir=resolve_project_path(data["dataset_dir"], config_path), |
| years=years, |
| variables=data["variables"], |
| task="pretrain", |
| input_steps=int(data["input_steps"]), |
| output_steps=int(data["output_steps"]), |
| normalize=bool(data["normalize"]), |
| ) |
|
|
|
|
| def validate_architecture_reference( |
| model: torch.nn.Module, |
| reference_path: Path, |
| config_path: Path, |
| source_root: Path | None, |
| ) -> None: |
| reference_model = build_model(load_config(config_path)) |
| load_checkpoint(reference_model, reference_path, source_root=source_root) |
| model_state = model.state_dict() |
| reference_state = reference_model.state_dict() |
| if set(model_state) != set(reference_state): |
| raise RuntimeError("Model parameter names do not match the official architecture reference.") |
| mismatches = { |
| key: (tuple(model_state[key].shape), tuple(reference_state[key].shape)) |
| for key in model_state |
| if tuple(model_state[key].shape) != tuple(reference_state[key].shape) |
| } |
| if mismatches: |
| raise RuntimeError(f"Model parameter shapes do not match the official architecture: {mismatches}") |
|
|
|
|
| def save_checkpoint( |
| path: Path, |
| model: torch.nn.Module, |
| optimizer: torch.optim.Optimizer, |
| epoch: int, |
| loss: float, |
| config: dict[str, Any], |
| ) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| torch.save( |
| { |
| "model": model.state_dict(), |
| "optimizer": optimizer.state_dict(), |
| "epoch": epoch, |
| "loss": loss, |
| "config": config, |
| }, |
| path, |
| ) |
|
|
|
|
| def resolve_device(requested: str) -> torch.device: |
| if requested == "auto": |
| requested = "cuda" if torch.cuda.is_available() else "cpu" |
| device = torch.device(requested) |
| if device.type == "cuda" and not torch.cuda.is_available(): |
| raise RuntimeError("The requested cuda device is unavailable. Use --device cpu or check the DCU runtime.") |
| return device |
|
|
|
|
| def device_summary(device: torch.device) -> str: |
| if device.type != "cuda": |
| return f"device={device}" |
| backend = f"HIP {torch.version.hip}" if torch.version.hip else f"CUDA {torch.version.cuda}" |
| return f"device={device} name={torch.cuda.get_device_name(device)} backend={backend}" |
|
|
|
|
| def memory_summary(device: torch.device) -> str: |
| if device.type != "cuda": |
| return "memory=unavailable(cpu)" |
| allocated = torch.cuda.memory_allocated(device) / (1024**3) |
| reserved = torch.cuda.memory_reserved(device) / (1024**3) |
| peak = torch.cuda.max_memory_allocated(device) / (1024**3) |
| return f"memory_allocated={allocated:.2f}GB memory_reserved={reserved:.2f}GB peak_allocated={peak:.2f}GB" |
|
|
|
|
| def setup_distributed(requested_device: str) -> tuple[torch.device, int, int, int]: |
| world_size = int(os.environ.get("WORLD_SIZE", "1")) |
| rank = int(os.environ.get("RANK", "0")) |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) |
| if world_size == 1: |
| return resolve_device(requested_device), rank, world_size, local_rank |
|
|
| if requested_device == "cpu": |
| device = torch.device("cpu") |
| backend = "gloo" |
| else: |
| if not torch.cuda.is_available(): |
| raise RuntimeError("torchrun requires one accessible DCU/GPU per local rank.") |
| if local_rank >= torch.cuda.device_count(): |
| raise RuntimeError( |
| f"LOCAL_RANK={local_rank} exceeds the {torch.cuda.device_count()} visible accelerator devices." |
| ) |
| torch.cuda.set_device(local_rank) |
| device = torch.device("cuda", local_rank) |
| backend = "nccl" |
| dist.init_process_group(backend=backend, init_method="env://") |
| return device, rank, world_size, local_rank |
|
|
|
|
| def rank_print(message: str, rank: int) -> None: |
| if rank == 0: |
| print(message, flush=True) |
|
|
|
|
| def all_ranks_finite(value: torch.Tensor, world_size: int) -> bool: |
| finite = torch.isfinite(value).to(dtype=torch.int32) |
| if world_size > 1: |
| dist.all_reduce(finite, op=dist.ReduceOp.MIN) |
| return bool(finite.item()) |
|
|
|
|
| def unwrap_model(model: torch.nn.Module) -> torch.nn.Module: |
| return model.module if isinstance(model, DistributedDataParallel) else model |
|
|
|
|
| def train_one_epoch( |
| model: torch.nn.Module, |
| loader: DataLoader, |
| optimizer: torch.optim.Optimizer, |
| mask_ratio: float, |
| device: torch.device, |
| epoch: int, |
| log_interval: int, |
| rank: int, |
| world_size: int, |
| ) -> float: |
| model.train() |
| total_loss = 0.0 |
| batches = 0 |
| total_batches = len(loader) |
| if device.type == "cuda": |
| torch.cuda.reset_peak_memory_stats(device) |
| for batch_index, (inputs, _, _, _, _) in enumerate(loader, start=1): |
| inputs = inputs.to(device) |
| output = model(inputs, mask_ratio=mask_ratio) |
| if not all_ranks_finite(output.loss.detach(), world_size): |
| raise FloatingPointError("Training loss is NaN or Inf on at least one rank.") |
| optimizer.zero_grad(set_to_none=True) |
| output.loss.backward() |
| optimizer.step() |
| total_loss += float(output.loss.detach().cpu()) |
| batches += 1 |
| if batch_index == 1 or batch_index % log_interval == 0 or batch_index == total_batches: |
| display_loss = output.loss.detach() |
| if world_size > 1: |
| dist.all_reduce(display_loss, op=dist.ReduceOp.SUM) |
| display_loss /= world_size |
| rank_print( |
| f"epoch={epoch} batch={batch_index}/{total_batches} loss={float(display_loss):.6f} " |
| f"{memory_summary(device)}", |
| rank, |
| ) |
| if batches == 0: |
| raise ValueError("Training DataLoader produced no batches.") |
| totals = torch.tensor([total_loss, batches], dtype=torch.float64, device=device) |
| if world_size > 1: |
| dist.all_reduce(totals, op=dist.ReduceOp.SUM) |
| return float((totals[0] / totals[1]).cpu()) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Train W-MAE using the configured ERA5 adapter.") |
| parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "conf" / "config.yaml") |
| parser.add_argument("--epochs", type=int, default=1) |
| parser.add_argument("--learning-rate", type=float, default=1e-4) |
| parser.add_argument( |
| "--device", |
| default="auto", |
| help="Training device: auto (default), cuda (AMD DCU through HIP), or cpu.", |
| ) |
| parser.add_argument("--checkpoint", type=Path, default=None) |
| parser.add_argument("--checkpoint-source-root", type=Path, default=None) |
| parser.add_argument("--non-strict-checkpoint", action="store_true") |
| parser.add_argument( |
| "--output-checkpoint", |
| type=Path, |
| default=PROJECT_ROOT / "data" / "checkpoint" / "model_bak.pth", |
| ) |
| parser.add_argument("--architecture-reference-checkpoint", type=Path, default=None) |
| parser.add_argument("--log-interval", type=int, default=1) |
| args = parser.parse_args() |
| if args.epochs <= 0: |
| raise ValueError("--epochs must be positive.") |
| if args.log_interval <= 0: |
| raise ValueError("--log-interval must be positive.") |
|
|
| config_path = args.config.resolve() |
| device, rank, world_size, local_rank = setup_distributed(args.device) |
| rank_print( |
| f"training started: config={config_path} epochs={args.epochs} device={args.device}", rank |
| ) |
| rank_print( |
| f"distributed: enabled={world_size > 1} world_size={world_size} " |
| f"rank={rank} local_rank={local_rank}", |
| rank, |
| ) |
| config = load_config(config_path) |
| validate_config(config) |
| data_config = config["data"] |
| print(f"rank={rank} local_rank={local_rank} runtime: {device_summary(device)}", flush=True) |
| rank_print("building model", rank) |
| model = build_model(config).to(device) |
| print( |
| f"rank={rank} local_rank={local_rank} model_parameters_device=" |
| f"{next(model.parameters()).device}", |
| flush=True, |
| ) |
|
|
| if args.checkpoint is not None: |
| checkpoint_path = resolve_project_path(args.checkpoint, config_path) |
| source_root = ( |
| resolve_project_path(args.checkpoint_source_root, config_path) |
| if args.checkpoint_source_root |
| else None |
| ) |
| report = load_checkpoint( |
| model, |
| checkpoint_path, |
| strict=not args.non_strict_checkpoint, |
| map_location=device, |
| source_root=source_root, |
| ) |
| rank_print(f"loaded checkpoint: {report}", rank) |
| else: |
| rank_print("training from randomly initialized weights", rank) |
|
|
| if world_size > 1: |
| if device.type == "cuda": |
| model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank) |
| else: |
| model = DistributedDataParallel(model) |
| rank_print("building dataset", rank) |
| dataset = build_dataset(config, config_path, "train") |
| sampler = DistributedSampler(dataset, shuffle=True) if world_size > 1 else None |
| loader = DataLoader( |
| dataset, |
| batch_size=int(data_config["batch_size"]), |
| shuffle=sampler is None, |
| sampler=sampler, |
| num_workers=int(data_config["num_workers"]), |
| ) |
| rank_print( |
| f"dataset ready: samples={len(dataset)} batches_per_rank={len(loader)} " |
| f"batch_size_per_rank={data_config['batch_size']} " |
| f"global_batch_size={int(data_config['batch_size']) * world_size}", |
| rank, |
| ) |
| optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate) |
| for epoch in range(args.epochs): |
| if sampler is not None: |
| sampler.set_epoch(epoch) |
| rank_print(f"epoch={epoch + 1}/{args.epochs} started", rank) |
| loss = train_one_epoch( |
| model, |
| loader, |
| optimizer, |
| float(config["model"]["mask_ratio"]), |
| device, |
| epoch + 1, |
| args.log_interval, |
| rank, |
| world_size, |
| ) |
| rank_print(f"epoch={epoch + 1}/{args.epochs} finished loss={loss:.6f}", rank) |
|
|
| if world_size > 1: |
| dist.barrier() |
|
|
| if rank != 0: |
| return |
|
|
| checkpoint_model = unwrap_model(model) |
|
|
| reference_path = ( |
| resolve_project_path(args.architecture_reference_checkpoint, config_path) |
| if args.architecture_reference_checkpoint |
| else None |
| ) |
| reference_source_root = ( |
| resolve_project_path(args.checkpoint_source_root, config_path) |
| if args.checkpoint_source_root |
| else None |
| ) |
| if reference_path is not None: |
| print(f"checking architecture against {reference_path}", flush=True) |
| validate_architecture_reference(checkpoint_model, reference_path, config_path, reference_source_root) |
| print("architecture reference check passed", flush=True) |
| else: |
| print("architecture reference check skipped: no reference checkpoint supplied", flush=True) |
| output_path = resolve_project_path(args.output_checkpoint, config_path) |
| save_checkpoint(output_path, checkpoint_model, optimizer, args.epochs, loss, config) |
| print(f"checkpoint saved: {output_path}", flush=True) |
| print("training completed successfully", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| try: |
| main() |
| except Exception: |
| traceback.print_exc() |
| raise |
| finally: |
| if dist.is_available() and dist.is_initialized(): |
| dist.destroy_process_group() |
|
|
