| """DDP-capable FuXi-DA training with analysis and frozen-proxy forecast supervision.""" |
| import argparse |
| import json |
| import math |
| import os |
| from pathlib import Path |
|
|
| import torch |
| import torch.distributed as dist |
| from torch.nn.parallel import DistributedDataParallel |
| from torch.utils.data import DataLoader, DistributedSampler |
|
|
| import sys |
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| import yaml |
| from model.fuxi_da import CompactForecastProxy, FuXiDA, ProceduralTileDataset |
|
|
|
|
| def latitude_weighted_l1(prediction, target, latitude): |
| weight = torch.cos(torch.deg2rad(latitude)).clamp_min(0) |
| weight = weight * weight.shape[-1] / weight.sum(dim=-1, keepdim=True) |
| return ((prediction - target).abs() * weight[:, None, :, None]).mean() |
|
|
|
|
| def learning_rate(step, warmup, total, peak): |
| if step < warmup: |
| return 1e-8 + (peak - 1e-8) * step / max(1, warmup) |
| progress = (step - warmup) / max(1, total - warmup) |
| return peak * 0.5 * (1.0 + math.cos(math.pi * progress)) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", default="conf/config.yaml") |
| parser.add_argument("--iterations", type=int) |
| parser.add_argument("--forecast-steps", type=int) |
| parser.add_argument("--output-dir") |
| args = parser.parse_args() |
| cfg = yaml.safe_load((ROOT / args.config).read_text()) |
| total = args.iterations or cfg["train"]["iterations"] |
| forecast_steps = args.forecast_steps or cfg["train"]["forecast_steps"] |
|
|
| world_size = int(os.environ.get("WORLD_SIZE", "1")) |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) |
| use_cuda = torch.cuda.is_available() and torch.cuda.device_count() >= world_size and cfg["runtime"]["device"] != "cpu" |
| if world_size > 1: |
| dist.init_process_group("nccl" if use_cuda else "gloo") |
| device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu") |
| if use_cuda: |
| torch.cuda.set_device(device) |
| torch.manual_seed(cfg["seed"] + local_rank) |
|
|
| dataset = ProceduralTileDataset(cfg["data"]["tile_ids"], max(2, total), cfg["data"]["tile_size"], cfg["data"]["missing_probability"]) |
| sampler = DistributedSampler(dataset, shuffle=True) if world_size > 1 else None |
| loader = DataLoader(dataset, batch_size=cfg["train"]["batch_size"], sampler=sampler, shuffle=sampler is None, num_workers=0) |
| model = FuXiDA(cfg["model"]["base_channels"]).to(device) |
| proxy = CompactForecastProxy().to(device) |
| if world_size > 1: |
| model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) |
| optimizer = torch.optim.AdamW(model.parameters(), lr=cfg["train"]["learning_rate"], betas=(0.9, 0.999), weight_decay=cfg["train"]["weight_decay"]) |
|
|
| model.train() |
| iterator = iter(loader) |
| for step in range(total): |
| try: |
| batch = next(iterator) |
| except StopIteration: |
| if sampler is not None: |
| sampler.set_epoch(step) |
| iterator = iter(loader) |
| batch = next(iterator) |
| background, obs, target = (batch[key].to(device) for key in ("background", "obs", "target")) |
| latitude = batch["latitude"].to(device) |
| analysis = model(background, obs) |
| analysis_loss = latitude_weighted_l1(analysis, target, latitude) |
| state, forecast_loss = analysis, analysis_loss.new_zeros(()) |
| for lead in range(forecast_steps): |
| state = proxy(state) |
| forecast_loss = forecast_loss + latitude_weighted_l1(state, batch["forecast_targets"][:, lead].to(device), latitude) |
| loss = analysis_loss + forecast_loss / forecast_steps |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| optimizer.step() |
| lr = learning_rate(step + 1, cfg["train"]["warmup_steps"], total, cfg["train"]["learning_rate"]) |
| for group in optimizer.param_groups: |
| group["lr"] = lr |
| if local_rank == 0 and (step == 0 or (step + 1) % 100 == 0 or step + 1 == total): |
| print(json.dumps({"step": step + 1, "loss": loss.item(), "analysis_l1": analysis_loss.item(), "lr": lr})) |
|
|
| if local_rank == 0: |
| checkpoint = ROOT / cfg["paths"]["checkpoint"]; checkpoint.parent.mkdir(parents=True, exist_ok=True) |
| bare_model = model.module if isinstance(model, DistributedDataParallel) else model |
| torch.save({"model": bare_model.state_dict(), "model_config": cfg["model"], "format_version": cfg["data"]["format_version"]}, checkpoint) |
| metrics = ROOT / cfg["paths"]["training_metrics"]; metrics.parent.mkdir(parents=True, exist_ok=True) |
| metrics.write_text(json.dumps({"iterations": total, "final_loss": float(loss), "world_size": world_size}, indent=2)) |
| if world_size > 1: |
| dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
