Publish FuXi-DA reproduction

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text

conf/config.yaml

@@ -0,0 +1,33 @@
+seed: 2025
+data:
+  format_version: fuxi_da_tiles_v1
+  manifest: data/fuxi_da_manifest.json
+  tile_size: 32
+  tile_ids: [0, 37, 181, 399]
+  missing_probability: 0.2
+model:
+  base_channels: 8
+train:
+  iterations: 2
+  forecast_steps: 2
+  batch_size: 1
+  learning_rate: 0.002
+  warmup_steps: 1
+  weight_decay: 0.00001
+runtime:
+  device: cpu
+paths:
+  checkpoint: result/checkpoints/fuxi_da.pt
+  training_metrics: result/training/metrics.json
+  predictions: result/output/predictions.npz
+  evaluation: result/evaluation/metrics.json
+  figure: result/evaluation/comparison.png
+paper_model:
+  background_shape: [70, 721, 1440]
+  observation_shape: [8, 15, 640, 640]
+  iterations: 6000
+  forecast_steps: 10
+  batch_size_per_gpu: 1
+  gpus: 4
+  warmup_steps: 500
+  peak_learning_rate: 0.002

config.json

@@ -0,0 +1 @@
+{"model_name":"FuXi-DA","model_type":"fuxi_da","architectures":["FuXiDA"],"framework":"PyTorch","domain":"weather","task":"satellite-data-assimilation","implementation":{"entry_point":"model/fuxi_da.py","scope":"core-method full-variable tile-streamed engineering reproduction","train_script":"scripts/train.py","inference_script":"scripts/inference.py","evaluation_script":"scripts/result.py","synthetic_data_script":"scripts/fake_data.py"},"architecture":{"background_shape":[70,721,1440],"observation_shape":[8,15,640,640],"analysis_shape":[70,720,1440],"engineering_tile":[32,32],"core":"multi-branch U-Net and multi-scale fusion"},"data":{"datasets":["ERA5","Fengyun-4B AGRI"],"format_version":"fuxi_da_tiles_v1","variables":70,"observation_channels":15,"observation_frames":8,"synthetic":true},"configuration_sources":["conf/config.yaml","model/fuxi_da.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/fuxi_da.py

@@ -0,0 +1,144 @@
+"""Independent PyTorch implementation of a compact multi-branch FuXi-DA U-Net."""
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+
+BACKGROUND_RAW_SHAPE = (70, 721, 1440)
+BACKGROUND_MODEL_SHAPE = (70, 720, 1440)
+AGRI_RAW_SHAPE = (8, 15, 640, 640)
+ANALYSIS_SHAPE = (70, 720, 1440)
+FOOTPRINT_ORIGIN = (40, 400)
+
+
+def tile_origin(tile_id, tile_size=32):
+    columns = 640 // tile_size
+    return FOOTPRINT_ORIGIN[0] + (tile_id // columns) * tile_size, FOOTPRINT_ORIGIN[1] + (tile_id % columns) * tile_size
+
+
+def _coords(y0, x0, size):
+    y = torch.arange(y0, y0 + size).float(); x = torch.arange(x0, x0 + size).float()
+    lat = 90 - (y + .5) * 180 / 720; lon = -180 + (x + .5) * 360 / 1440
+    return torch.meshgrid(torch.deg2rad(lat), torch.deg2rad(lon), indexing="ij")
+
+
+def procedural_state(y0, x0, size, sample, lead=0):
+    lat, lon = _coords(y0, x0, size); channel = torch.arange(70).float()[:, None, None]
+    phase = .07 * sample + .11 * lead + channel * .031
+    return torch.sin((1 + channel % 4) * lat + phase) * torch.cos((1 + channel % 5) * lon - phase) + .3 * torch.sin(3 * lon + 2 * lat + phase)
+
+
+def make_sample(tile_id, sample, size=32, missing_probability=.2):
+    y0, x0 = tile_origin(tile_id, size); target = procedural_state(y0, x0, size, sample)
+    lat, lon = _coords(y0, x0, size); channels = torch.arange(70).float()[:, None, None]
+    error = .12 * torch.sin(2 * lon + channels * .09) + .05 * torch.cos(3 * lat)
+    background = target + error
+    times = torch.arange(8).float()[:, None, None, None]; variables = torch.arange(15).long()
+    obs = target[(variables * 4) % 70][None] + .35 * background[(variables * 3) % 70][None] + .02 * times
+    generator = torch.Generator().manual_seed(100003 * sample + tile_id)
+    missing = torch.rand((8, 1, size, size), generator=generator) < missing_probability
+    obs = obs.masked_fill(missing.expand_as(obs), 0).reshape(120, size, size)
+    forecasts = torch.stack([procedural_state(y0, x0, size, sample, i + 1) for i in range(10)])
+    return {"background": background, "obs": obs, "target": target, "correction": background - error, "forecast_targets": forecasts,
+            "latitude": torch.rad2deg(lat[:, 0]), "origin": torch.tensor([y0, x0]), "tile_id": torch.tensor(tile_id)}
+
+
+class ProceduralTileDataset(Dataset):
+    def __init__(self, tile_ids, samples_per_tile=2, tile_size=32, missing_probability=.2):
+        self.tile_ids, self.samples_per_tile, self.tile_size, self.missing_probability = tuple(tile_ids), samples_per_tile, tile_size, missing_probability
+    def __len__(self): return len(self.tile_ids) * self.samples_per_tile
+    def __getitem__(self, index):
+        return make_sample(self.tile_ids[index % len(self.tile_ids)], index // len(self.tile_ids), self.tile_size, self.missing_probability)
+
+
+class ChannelLayerNorm(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.norm = nn.LayerNorm(channels)
+
+    def forward(self, x):
+        return self.norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+
+class Down(nn.Sequential):
+    def __init__(self, input_channels, output_channels):
+        super().__init__(
+            nn.Conv2d(input_channels, output_channels, 2, stride=2),
+            ChannelLayerNorm(output_channels),
+            nn.SiLU(),
+            nn.Conv2d(output_channels, output_channels, 3, padding=1),
+        )
+
+
+class Up(nn.Sequential):
+    def __init__(self, input_channels, output_channels):
+        super().__init__(
+            nn.Conv2d(input_channels, input_channels, 3, padding=1),
+            ChannelLayerNorm(input_channels),
+            nn.SiLU(),
+            nn.Conv2d(input_channels, output_channels * 4, 3, padding=1),
+            nn.PixelShuffle(2),
+        )
+
+
+class Stem(nn.Sequential):
+    def __init__(self, input_channels, channels):
+        super().__init__(nn.Conv2d(input_channels, channels, 3, padding=1), ChannelLayerNorm(channels), nn.SiLU())
+
+
+class Fusion(nn.Module):
+    """Interact modalities and emit increment, observation bias, and corrected condition."""
+    def __init__(self, channels):
+        super().__init__()
+        self.mix = nn.Sequential(
+            nn.Conv2d(channels * 3, channels, 3, padding=1), ChannelLayerNorm(channels), nn.SiLU(),
+            nn.Conv2d(channels, channels * 3, 3, padding=1),
+        )
+
+    def forward(self, background, observation, condition):
+        increment, obs_bias, corrected_condition = self.mix(torch.cat((background, observation, condition), dim=1)).chunk(3, dim=1)
+        return background + increment, observation + obs_bias, condition + corrected_condition
+
+
+class FuXiDA(nn.Module):
+    def __init__(self, base_channels=8):
+        super().__init__()
+        c = base_channels
+        self.background_stem, self.obs_stem, self.condition_stem = Stem(70, c), Stem(120, c), Stem(190, c)
+        self.fuse0 = Fusion(c)
+        self.background_down1, self.obs_down1, self.condition_down1 = Down(c, 2*c), Down(c, 2*c), Down(c, 2*c)
+        self.fuse1 = Fusion(2*c)
+        self.background_down2, self.obs_down2, self.condition_down2 = Down(2*c, 4*c), Down(2*c, 4*c), Down(2*c, 4*c)
+        self.fuse2 = Fusion(4*c)
+        self.up1 = Up(4*c, 2*c)
+        self.decode1 = nn.Sequential(nn.Conv2d(4*c, 2*c, 3, padding=1), ChannelLayerNorm(2*c), nn.SiLU())
+        self.up0 = Up(2*c, c)
+        self.decode0 = nn.Sequential(nn.Conv2d(2*c, c, 3, padding=1), ChannelLayerNorm(c), nn.SiLU())
+        self.increment_head = nn.Conv2d(c, 70, 3, padding=1)
+
+    def forward(self, background, observation):
+        condition = torch.cat((background, observation), dim=1)
+        b0, o0, c0 = self.fuse0(self.background_stem(background), self.obs_stem(observation), self.condition_stem(condition))
+        b1, o1, c1 = self.fuse1(self.background_down1(b0), self.obs_down1(o0), self.condition_down1(c0))
+        b2, _, _ = self.fuse2(self.background_down2(b1), self.obs_down2(o1), self.condition_down2(c1))
+        decoded = self.decode1(torch.cat((self.up1(b2), b1), dim=1))
+        decoded = self.decode0(torch.cat((self.up0(decoded), b0), dim=1))
+        increment = self.increment_head(decoded)
+        return background + increment
+
+
+class CompactForecastProxy(nn.Module):
+    """Frozen differentiable forecast surrogate; gradients only update FuXi-DA."""
+    def __init__(self):
+        super().__init__()
+        self.depthwise = nn.Conv2d(70, 70, 3, padding=1, groups=70, bias=False)
+        with torch.no_grad():
+            self.depthwise.weight.fill_(1.0 / 9.0)
+        self.requires_grad_(False)
+        self.eval()
+
+    def forward(self, state):
+        return 0.985 * self.depthwise(state) + 0.015 * torch.roll(state, shifts=1, dims=-1)
+
+    def train(self, mode=True):
+        return super().train(False)

scripts/fake_data.py

@@ -0,0 +1,8 @@
+from pathlib import Path
+import json,sys
+import yaml
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.fuxi_da import *
+c=yaml.safe_load((ROOT/"conf/config.yaml").read_text()); sample=make_sample(c["data"]["tile_ids"][0],0,c["data"]["tile_size"],c["data"]["missing_probability"])
+assert BACKGROUND_RAW_SHAPE==(70,721,1440) and sample["background"].shape==(70,c["data"]["tile_size"],c["data"]["tile_size"]) and sample["obs"].shape[0]==120
+path=ROOT/c["data"]["manifest"];path.parent.mkdir(parents=True,exist_ok=True);path.write_text(json.dumps({"format_version":c["data"]["format_version"],"background_raw_shape":[70,721,1440],"background_model_shape":[70,720,1440],"agri_raw_shape":[8,15,640,640],"observation_model_shape":[120,640,640],"condition_shape":[190,640,640],"analysis_shape":[70,720,1440],"tile_ids":c["data"]["tile_ids"],"tile_size":c["data"]["tile_size"],"is_complete_global":False},indent=2));print(path)

scripts/inference.py

@@ -0,0 +1,9 @@
+from pathlib import Path
+import sys,yaml,numpy as np,torch
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.fuxi_da import FuXiDA,make_sample
+c=yaml.safe_load((ROOT/"conf/config.yaml").read_text());ck=torch.load(ROOT/c["paths"]["checkpoint"],map_location="cpu",weights_only=True);assert ck["format_version"]==c["data"]["format_version"];m=FuXiDA(**ck["model_config"]);m.load_state_dict(ck["model"]);m.eval();rows=[]
+with torch.no_grad():
+ for tile in c["data"]["tile_ids"]:
+  s=make_sample(tile,20,c["data"]["tile_size"],c["data"]["missing_probability"]);p=m(s["background"][None],s["obs"][None])[0];rows.append((s,p))
+path=ROOT/c["paths"]["predictions"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,background=np.stack([r[0]["background"].numpy() for r in rows]),observations=np.stack([r[0]["obs"].numpy() for r in rows]),target=np.stack([r[0]["target"].numpy() for r in rows]),prediction=np.stack([r[1].numpy() for r in rows]),latitude=np.stack([r[0]["latitude"].numpy() for r in rows]),origins=np.stack([r[0]["origin"].numpy() for r in rows]),tile_ids=np.array(c["data"]["tile_ids"]),format_version=np.array(c["data"]["format_version"]),is_complete_global=np.bool_(False));print(path)

scripts/result.py

@@ -0,0 +1,85 @@
+"""Tile diagnostics for analysis, forecasts, variable groups, and localization."""
+import argparse
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import sys, yaml
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from model.fuxi_da import CompactForecastProxy, FuXiDA, make_sample
+
+GROUPS = {"Z": slice(0, 13), "T": slice(13, 26), "U": slice(26, 39), "V": slice(39, 52), "R": slice(52, 65), "surface": slice(65, 70)}
+
+
+def weighted_rmse(prediction, target, latitude):
+    weight = torch.cos(torch.deg2rad(latitude)).clamp_min(0)
+    weight = weight / weight.sum()
+    return torch.sqrt(((prediction - target).square() * weight[None, :, None]).sum(dim=(-2, -1)) / prediction.shape[-1]).mean().item()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--output")
+    parser.add_argument("--metrics")
+    args = parser.parse_args()
+    cfg = yaml.safe_load((ROOT / args.config).read_text())
+    model = FuXiDA(cfg["model"]["base_channels"])
+    checkpoint_path = Path(args.checkpoint) if args.checkpoint else ROOT / cfg["paths"]["checkpoint"]
+    if checkpoint_path.exists():
+        model.load_state_dict(torch.load(checkpoint_path, map_location="cpu", weights_only=True)["model"])
+    model.eval()
+    proxy = CompactForecastProxy()
+
+    rows = []
+    with torch.no_grad():
+        for tile_id in cfg["data"]["tile_ids"]:
+            sample = make_sample(tile_id, 20, cfg["data"]["tile_size"], cfg["data"]["missing_probability"])
+            analysis = model(sample["background"][None], sample["obs"][None])[0]
+            row = {"tile_id": tile_id, "origin": sample["origin"].tolist(), "background": weighted_rmse(sample["background"], sample["target"], sample["latitude"]), "correction": weighted_rmse(sample["correction"], sample["target"], sample["latitude"]), "analysis": weighted_rmse(analysis, sample["target"], sample["latitude"])}
+            row["variable_groups"] = {name: weighted_rmse(analysis[index], sample["target"][index], sample["latitude"]) for name, index in GROUPS.items()}
+            state = analysis
+            row["forecast_steps"] = []
+            for lead in range(cfg["train"]["forecast_steps"]):
+                state = proxy(state[None])[0]
+                row["forecast_steps"].append(weighted_rmse(state, sample["forecast_targets"][lead], sample["latitude"]))
+            rows.append(row)
+
+        sample = make_sample(cfg["data"]["tile_ids"][0], 21, cfg["data"]["tile_size"], 0.0)
+        base = model(sample["background"][None], sample["obs"][None])[0]
+        perturbed_obs = sample["obs"].clone()
+        center = cfg["data"]["tile_size"] // 2
+        perturbed_obs[9 - 8, center, center] += 1.0
+        response = (model(sample["background"][None], perturbed_obs[None])[0] - base).square().sum(0)
+        yy, xx = torch.meshgrid(torch.arange(cfg["data"]["tile_size"]), torch.arange(cfg["data"]["tile_size"]), indexing="ij")
+        radius = torch.sqrt((yy - center).float().square() + (xx - center).float().square())
+        total_energy = response.sum().clamp_min(1e-12)
+        localization = {"perturbation": "AGRI channel 9 +1 K at tile center", "energy_weighted_radius_gridpoints": (response * radius).sum().div(total_energy).item(), "energy_within_radius_4": response[radius <= 4].sum().div(total_energy).item()}
+
+    metrics = {"coverage_complete": False, "claim": "diagnostic metrics on selected aligned tiles; not complete global scores", "tiles": rows, "increment_localization": localization}
+    metrics_path = Path(args.metrics) if args.metrics else ROOT / cfg["paths"]["evaluation"]
+    metrics_path.parent.mkdir(parents=True, exist_ok=True); metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
+    labels = ["background", "correction", "analysis"]
+    values = [np.mean([row[label] for row in rows]) for label in labels]
+    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
+    axes[0].bar(labels, values, color=["#596780", "#d69b36", "#287f71"])
+    axes[0].set_ylabel("Latitude-weighted RMSE")
+    axes[0].set_title("Selected aligned tiles (incomplete coverage)")
+    group_values = [np.mean([row["variable_groups"][name] for row in rows]) for name in GROUPS]
+    axes[1].bar(list(GROUPS), group_values, color="#287f71")
+    axes[1].set_title("Analysis RMSE by variable group")
+    fig.suptitle("FuXi-DA procedural protocol diagnostics")
+    fig.tight_layout()
+    output_path = Path(args.output) if args.output else ROOT / cfg["paths"]["figure"]
+    output_path.parent.mkdir(parents=True, exist_ok=True); fig.savefig(output_path, dpi=160)
+    plt.close(fig)
+    print(json.dumps(metrics, indent=2))
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,102 @@
+"""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()