Add files using upload-large-folder tool

README.md

@@ -97,17 +97,22 @@ bash Train/run_download_prepare.sh
 
 ## Training
 
-The default runnable configuration uses a direct batch size of 256 on one GPU,
-no gradient accumulation, 15 epochs, and saves checkpoints every 5 epochs:
+The default runnable configuration uses 8 GPUs, a global batch size of 256
+(32 samples per GPU), no gradient accumulation, 15 epochs, and saves
+checkpoints every 5 epochs:
 
 ```bash
-bash Train/run_train_local.sh
+bash Train/run_train_local.sh --num-gpus 8
 ```
 
+The GPU count is a launcher parameter. For example, use `--num-gpus 1`,
+`--num-gpus 4`, or `--num-gpus 8`. It defaults to 8 and may also be set with
+`NUM_GPUS`. The global batch size must be divisible by the GPU count.
+
 For a new machine, the complete download-to-training sequence can be started with:
 
 ```bash
-bash Train/run_all_local.sh
+bash Train/run_all_local.sh --num-gpus 8
 ```
 
 The training entrypoint itself is independent of a scheduler:
@@ -117,8 +122,10 @@ python Train/training/train.py --config Train/configs/stead_empty8_local.yaml
 ```
 
 The transfer-learning configuration in `Train/Transfer_and_Test/configs/config.yaml`
-also uses a training batch size of 256. Generation uses a batch size of 256 in
-both local and transfer-learning launchers.
+also uses a global training batch size of 256. Generation uses a global batch
+size of 256 in both local and transfer-learning launchers. Training and
+generation distribute their work across the requested GPUs and only the main
+process writes the final checkpoint or merged HDF5 output.
 
 If another machine cannot fit 256 waveforms in GPU memory, reduce the configured
 batch size before running. If distributed training is required, launch this same
@@ -133,5 +140,11 @@ After epoch 15 finishes:
 bash Train/run_generate_100.sh
 ```
 
+This also defaults to 8 GPUs. To select a different count:
+
+```bash
+bash Train/run_generate_100.sh --num-gpus 4
+```
+
 This performs ancestral DDPM sampling with 25 steps and seed 2026. The output
 is `Train/results/stead_empty8_local/generated_100/generated_ddpm25_seed2026.h5`.

Transfer_and_Test/README.md

@@ -15,10 +15,12 @@ Run from the `GenerationArticle` directory:
 
 ```bash
 conda activate DiT
-bash Train/Transfer_and_Test/run_transfer_and_test.sh
+bash Train/Transfer_and_Test/run_transfer_and_test.sh --num-gpus 8
 ```
 
-The launcher uses paths inside this directory, trains with batch size 256,
-saves checkpoints every 50 epochs, generates 100 DDPM samples in a batch of
-256 at each milestone, and evaluates them with
-the bundled EQT checkpoint using a 1.0-second tolerance.
+`--num-gpus` is configurable and defaults to 8. The launcher uses paths inside
+this directory, trains with global batch size 256 (32 per GPU with 8 GPUs),
+saves checkpoints every 50 epochs, and distributes each 256-sample generation
+batch across the selected GPUs before merging the shards. At each milestone it
+generates 100 DDPM samples and evaluates them with
+the bundled EQT checkpoint using the selected GPUs and a 1.0-second tolerance.

Transfer_and_Test/run_transfer_and_test.sh

@@ -17,6 +17,24 @@ EVAL_SCRIPT="${BUNDLE_ROOT}/src/eqt/evaluate_transfer_eqt_test_in_memory.py"
 EQT_CHECKPOINT="${BUNDLE_ROOT}/eqt/best_val_state_dict.pt"
 LOG_DIR="${RESULT_ROOT}/logs"
 MILESTONE_ROOT="${RESULT_ROOT}/milestones"
+NUM_GPUS="${NUM_GPUS:-8}"
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --num-gpus)
+      [[ $# -ge 2 ]] || { echo "--num-gpus requires a value" >&2; exit 2; }
+      NUM_GPUS="$2"
+      shift 2
+      ;;
+    --num-gpus=*) NUM_GPUS="${1#*=}"; shift ;;
+    *)
+      echo "Unknown argument: $1" >&2
+      echo "Usage: bash Train/Transfer_and_Test/run_transfer_and_test.sh [--num-gpus N]" >&2
+      exit 2
+      ;;
+  esac
+done
+[[ "${NUM_GPUS}" =~ ^[1-9][0-9]*$ ]] || { echo "NUM_GPUS must be a positive integer" >&2; exit 2; }
 
 mkdir -p "${LOG_DIR}" "${MILESTONE_ROOT}"
 cd "${BUNDLE_ROOT}"
@@ -34,6 +52,17 @@ else
   }
 fi
 
+run_gpu_job() {
+  if [[ "${NUM_GPUS}" -eq 1 ]]; then
+    "${PYTHON_BIN}" -u "$@"
+  else
+    "${PYTHON_BIN}" -u -m torch.distributed.run \
+      --standalone \
+      --nproc_per_node="${NUM_GPUS}" \
+      "$@"
+  fi
+}
+
 for path in "${CONFIG}" "${TRAIN_SCRIPT}" "${GENERATE_SCRIPT}" \
   "${EVAL_SCRIPT}" "${EQT_CHECKPOINT}" \
   "${BUNDLE_ROOT}/data/processed/iquique/train_3000_two_conditions.h5" \
@@ -52,8 +81,13 @@ print(yaml.safe_load(open(sys.argv[1], encoding="utf-8"))["training"]["global_ba
 PY
 )"
 echo "[INFO] train batch=${TRAIN_BATCH}"
-echo "[INFO] generation samples=100 generation batch=256 sampler=ddpm steps=100 cfg=1"
-echo "[INFO] EQT tolerance=100 samples (1.0 s)"
+if (( TRAIN_BATCH % NUM_GPUS != 0 )); then
+  echo "[FAIL] global train batch ${TRAIN_BATCH} is not divisible by NUM_GPUS=${NUM_GPUS}" >&2
+  exit 2
+fi
+echo "[INFO] num_gpus=${NUM_GPUS} train_global_batch=${TRAIN_BATCH} train_per_gpu_batch=$((TRAIN_BATCH / NUM_GPUS))"
+echo "[INFO] generation samples=100 generation_global_batch=256 sampler=ddpm steps=100 cfg=1"
+echo "[INFO] EQT num_gpus=${NUM_GPUS} tolerance=100 samples (1.0 s)"
 nvidia-smi || true
 
 FINAL_EPOCH="${TARGET_EPOCH:-2140}"
@@ -71,7 +105,7 @@ for epoch in "${MILESTONES[@]}"; do
   if [[ ! -f "${checkpoint}" ]]; then
     if [[ "${epoch}" -eq 50 ]]; then
       echo "[train] target_epoch=${epoch}"
-      "${PYTHON_BIN}" -u "${TRAIN_SCRIPT}" --config "${CONFIG}" \
+      run_gpu_job "${TRAIN_SCRIPT}" --config "${CONFIG}" \
         --target-epoch "${epoch}" \
         2>&1 | tee "${LOG_DIR}/train_to_${epoch}.log"
     else
@@ -86,7 +120,7 @@ for epoch in "${MILESTONES[@]}"; do
         exit 1
       }
       echo "[train] resume=${previous_checkpoint} target_epoch=${epoch}"
-      "${PYTHON_BIN}" -u "${TRAIN_SCRIPT}" --config "${CONFIG}" \
+      run_gpu_job "${TRAIN_SCRIPT}" --config "${CONFIG}" \
         --resume "${previous_checkpoint}" --target-epoch "${epoch}" \
         2>&1 | tee "${LOG_DIR}/train_to_${epoch}.log"
     fi
@@ -103,7 +137,7 @@ for epoch in "${MILESTONES[@]}"; do
 
   if [[ ! -f "${generated}" ]]; then
     echo "[generate] epoch=${epoch}"
-    "${PYTHON_BIN}" -u "${GENERATE_SCRIPT}" \
+    run_gpu_job "${GENERATE_SCRIPT}" \
       --checkpoint "${checkpoint}" \
       --source-h5 "${BUNDLE_ROOT}/data/processed/iquique/train_3000_two_conditions.h5" \
       --output-h5 "${generated}" \
@@ -128,6 +162,7 @@ for epoch in "${MILESTONES[@]}"; do
       --tasks milestone \
       --extra-checkpoints "milestone=${EQT_CHECKPOINT}" \
       --device cuda \
+      --num-gpus "${NUM_GPUS}" \
       --batch-size 100 \
       --workers 0 \
       --threshold 0.1 \

Transfer_and_Test/src/eqt/evaluate_transfer_eqt_test_in_memory.py

@@ -93,6 +93,7 @@ def evaluate_one(
     batch_size: int,
     workers: int,
     device: torch.device,
+    num_gpus: int,
     threshold: float,
     peak_distance_samples: int,
     tolerance_samples: int,
@@ -123,6 +124,8 @@ def evaluate_one(
     model: EQTransformer = build_eqtransformer_for_state(waves.shape[-1], state).to(device)
     model.load_state_dict(state)
     model.eval()
+    if device.type == "cuda" and num_gpus > 1:
+        model = torch.nn.DataParallel(model, device_ids=list(range(num_gpus)))
 
     rows: list[dict[str, Any]] = []
     prob_chunks: list[np.ndarray] = []
@@ -227,6 +230,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--tasks", nargs="*", default=DEFAULT_TASKS)
     parser.add_argument("--extra-checkpoints", nargs="*", default=[], help="Additional task=checkpoint entries, e.g. original=work/EQT/eqtransformer_original_state_dict.pt")
     parser.add_argument("--device", default="cuda")
+    parser.add_argument("--num-gpus", type=int, default=1)
     parser.add_argument("--batch-size", type=int, default=1280)
     parser.add_argument("--workers", type=int, default=8)
     parser.add_argument("--threshold", type=float, default=0.1)
@@ -261,6 +265,15 @@ def main() -> None:
     print(f"Tasks ({len(args.tasks)}): {' '.join(args.tasks)}")
 
     device = torch.device(str(args.device) if torch.cuda.is_available() else "cpu")
+    num_gpus = int(args.num_gpus)
+    if num_gpus < 1:
+        raise ValueError("--num-gpus must be a positive integer")
+    if device.type == "cuda" and num_gpus > torch.cuda.device_count():
+        raise ValueError(
+            f"Requested {num_gpus} GPUs, but only {torch.cuda.device_count()} are visible"
+        )
+    if device.type != "cuda":
+        num_gpus = 1
     rows = []
     for task in args.tasks:
         checkpoint = resolve_checkpoint(task, run_root, extra_checkpoints)
@@ -274,6 +287,7 @@ def main() -> None:
             batch_size=int(args.batch_size),
             workers=int(args.workers),
             device=device,
+            num_gpus=num_gpus,
             threshold=float(args.threshold),
             peak_distance_samples=int(args.peak_distance_samples),
             tolerance_samples=int(args.tolerance_samples),

Transfer_and_Test/src/generate_controlnet_two_condition.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import argparse
 import csv
+import os
 import random
 import sys
 from pathlib import Path
@@ -9,6 +10,7 @@ from pathlib import Path
 import h5py
 import numpy as np
 import torch
+import torch.distributed as dist
 from tqdm import tqdm
 
 # This generator is bundled directly under <bundle>/src on the server.
@@ -19,6 +21,19 @@ from diffusion import create_diffusion
 from models_controlnet_two_condition import TwoConditionControlNetSWaG
 
 
+def distributed_context() -> tuple[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:
+        dist.init_process_group("nccl")
+    return rank, local_rank, world_size
+
+
+def shard_path(output_path: Path, rank: int) -> Path:
+    return output_path.with_name(f".{output_path.name}.rank{rank}.partial.h5")
+
+
 def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--checkpoint", type=Path, required=True)
@@ -41,10 +56,16 @@ def main() -> None:
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required for sampling")
 
-    random.seed(args.seed)
-    np.random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    torch.cuda.manual_seed_all(args.seed)
+    rank, local_rank, world_size = distributed_context()
+    device = torch.device("cuda", local_rank)
+    torch.cuda.set_device(device)
+    if args.batch_size % world_size:
+        raise ValueError(f"batch-size must be divisible by world size ({world_size})")
+    local_batch = args.batch_size // world_size
+    random.seed(args.seed + rank)
+    np.random.seed(args.seed + rank)
+    torch.manual_seed(args.seed + rank)
+    torch.cuda.manual_seed_all(args.seed + rank)
     rng = np.random.default_rng(args.seed)
 
     with h5py.File(args.source_h5, "r") as handle:
@@ -64,36 +85,40 @@ def main() -> None:
 
     checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
     config = checkpoint["config"]
-    model = TwoConditionControlNetSWaG(**config["model"]).cuda().eval()
+    model = TwoConditionControlNetSWaG(**config["model"]).to(device).eval()
     model.load_state_dict(checkpoint["model"], strict=True)
+    del checkpoint
     respacing = f"ddim{args.sampling_steps}" if args.sampler == "ddim" else str(args.sampling_steps)
     diffusion = create_diffusion(timestep_respacing=respacing, **config["diffusion"])
-    generator = torch.Generator(device="cuda").manual_seed(args.seed)
+    generator = torch.Generator(device=device).manual_seed(args.seed + rank)
 
-    args.output_h5.parent.mkdir(parents=True, exist_ok=True)
-    args.output_meta.parent.mkdir(parents=True, exist_ok=True)
-    partial = args.output_h5.with_name(args.output_h5.name + ".partial")
-    partial.unlink(missing_ok=True)
+    if rank == 0:
+        args.output_h5.parent.mkdir(parents=True, exist_ok=True)
+        args.output_meta.parent.mkdir(parents=True, exist_ok=True)
+        for shard_rank in range(world_size):
+            shard_path(args.output_h5, shard_rank).unlink(missing_ok=True)
+    if world_size > 1:
+        dist.barrier()
     shape = (args.num_samples, int(config["model"]["length"]), int(config["model"]["in_channels"]))
-    with h5py.File(partial, "w") as handle:
-        output = handle.create_dataset("data", shape=shape, dtype="float32", compression="gzip", shuffle=True)
-        handle.create_dataset("labels", data=labels, dtype="float32")
-        handle.attrs.update(
-            checkpoint=str(args.checkpoint.resolve()),
-            sampling_steps=args.sampling_steps,
-            sampler=args.sampler,
-            cfg_scale=args.cfg_scale,
-            seed=args.seed,
-        )
+    assigned = np.array_split(np.arange(args.num_samples, dtype=np.int64), world_size)[rank]
+    local_shape = (len(assigned), shape[1], shape[2])
+    with h5py.File(shard_path(args.output_h5, rank), "w") as handle:
+        output = handle.create_dataset("data", shape=local_shape, dtype="float32", compression="gzip", shuffle=True)
+        handle.create_dataset("indices", data=assigned, dtype="int64")
         with torch.no_grad():
-            for start in tqdm(range(0, args.num_samples, args.batch_size), desc="sampling"):
-                end = min(start + args.batch_size, args.num_samples)
-                batch_labels = torch.from_numpy(labels[start:end]).cuda()
+            for start in tqdm(
+                range(0, len(assigned), local_batch),
+                desc=f"sampling rank {rank}",
+                disable=rank != 0,
+            ):
+                end = min(start + local_batch, len(assigned))
+                global_positions = assigned[start:end]
+                batch_labels = torch.from_numpy(labels[global_positions]).to(device)
                 noise = torch.randn(
                     end - start,
                     shape[2],
                     shape[1],
-                    device="cuda",
+                    device=device,
                     generator=generator,
                 )
 
@@ -106,25 +131,51 @@ def main() -> None:
                         model_fn,
                         noise.shape,
                         noise=noise,
-                        device=torch.device("cuda"),
+                        device=device,
                         progress=False,
                         clip_denoised=False,
                     )
                 output[start:end] = samples.float().transpose(1, 2).cpu().numpy().astype(np.float32)
-    partial.replace(args.output_h5)
-
-    with args.output_meta.open("w", newline="", encoding="utf-8") as handle:
-        writer = csv.writer(handle)
-        writer.writerow(["generated_index", "source_test_index", "p_index", "s_index"])
-        writer.writerows(
-            (
-                i,
-                int(indices[i]),
-                float(labels[i, 0]),
-                float(labels[i, 1]),
+
+    if world_size > 1:
+        dist.barrier()
+    if rank == 0:
+        partial = args.output_h5.with_name(args.output_h5.name + ".partial")
+        partial.unlink(missing_ok=True)
+        with h5py.File(partial, "w") as handle:
+            output = handle.create_dataset("data", shape=shape, dtype="float32", compression="gzip", shuffle=True)
+            handle.create_dataset("labels", data=labels, dtype="float32")
+            handle.attrs.update(
+                checkpoint=str(args.checkpoint.resolve()),
+                sampling_steps=args.sampling_steps,
+                sampler=args.sampler,
+                cfg_scale=args.cfg_scale,
+                seed=args.seed,
+                world_size=world_size,
+            )
+            for shard_rank in range(world_size):
+                current = shard_path(args.output_h5, shard_rank)
+                with h5py.File(current, "r") as shard:
+                    positions = np.asarray(shard["indices"], dtype=np.int64)
+                    output[positions] = shard["data"][:]
+                current.unlink()
+        partial.replace(args.output_h5)
+
+        with args.output_meta.open("w", newline="", encoding="utf-8") as handle:
+            writer = csv.writer(handle)
+            writer.writerow(["generated_index", "source_test_index", "p_index", "s_index"])
+            writer.writerows(
+                (
+                    i,
+                    int(indices[i]),
+                    float(labels[i, 0]),
+                    float(labels[i, 1]),
+                )
+                for i in range(len(labels))
             )
-            for i in range(len(labels))
-        )
+    if world_size > 1:
+        dist.barrier()
+        dist.destroy_process_group()
 
 
 if __name__ == "__main__":

Transfer_and_Test/src/train_controlnet_two_condition.py

@@ -1,12 +1,16 @@
 from __future__ import annotations
 
 import argparse
+import os
 import random
 from pathlib import Path
 
 import numpy as np
 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
 from tqdm import tqdm
 
 from config import load_config, resolve_path
@@ -19,9 +23,19 @@ from models_controlnet_two_condition import TwoConditionControlNetSWaG
 ROOT = Path(__file__).resolve().parents[1]
 
 
+def distributed_context() -> tuple[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:
+        dist.init_process_group("nccl")
+    return rank, local_rank, world_size
+
+
 def save_checkpoint(path, model, optimizer, epoch, step, config, include_optimizer: bool):
+    source = model.module if isinstance(model, DistributedDataParallel) else model
     payload = {
-        "model": model.state_dict(),
+        "model": source.state_dict(),
         "epoch": epoch,
         "step": step,
         "config": config,
@@ -44,27 +58,39 @@ def main():
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required")
 
+    rank, local_rank, world_size = distributed_context()
+    device = torch.device("cuda", local_rank)
+    torch.cuda.set_device(device)
     config_path = Path(args.config).resolve()
     config = load_config(config_path)
     training = config["training"]
-    seed = int(training.get("seed", 0))
+    seed = int(training.get("seed", 0)) + rank
     random.seed(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
     output = resolve_path(ROOT, config["output"]["directory"])
     checkpoint_dir = output / "checkpoints"
-    checkpoint_dir.mkdir(parents=True, exist_ok=True)
-    (output / "config.yaml").write_text(config_path.read_text(encoding="utf-8"), encoding="utf-8")
+    if rank == 0:
+        checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        (output / "config.yaml").write_text(config_path.read_text(encoding="utf-8"), encoding="utf-8")
+    if world_size > 1:
+        dist.barrier()
 
     dataset = TwoConditionDataset(
         resolve_path(ROOT, config["data"]["train_h5"]),
         cache_in_memory=bool(training.get("cache_in_memory", False)),
     )
+    global_batch = int(training["global_batch_size"])
+    if global_batch % world_size:
+        raise ValueError(f"global_batch_size must be divisible by world_size ({world_size})")
+    local_batch = global_batch // world_size
+    sampler = DistributedSampler(dataset, shuffle=True, seed=int(training.get("seed", 0))) if world_size > 1 else None
     loader = DataLoader(
         dataset,
-        batch_size=int(training["global_batch_size"]),
-        shuffle=True,
+        batch_size=local_batch,
+        shuffle=sampler is None,
+        sampler=sampler,
         num_workers=int(training.get("num_workers", 0)),
         pin_memory=True,
         drop_last=True,
@@ -117,12 +143,16 @@ def main():
     for name, parameter in model.named_parameters():
         parameter.requires_grad_(name.startswith(train_prefixes))
     trainable = [parameter for parameter in model.parameters() if parameter.requires_grad]
-    print(f"resume_checkpoint={resume_checkpoint}")
-    print(f"control_depth={model.control_depth}")
-    print(f"trainable_parameters={sum(p.numel() for p in trainable):,}")
-    print(f"frozen_parameters={sum(p.numel() for p in model.parameters() if not p.requires_grad):,}")
-
-    model.cuda()
+    if rank == 0:
+        print(f"resume_checkpoint={resume_checkpoint}")
+        print(f"world_size={world_size} global_batch={global_batch} per_gpu_batch={local_batch}")
+        print(f"control_depth={model.control_depth}")
+        print(f"trainable_parameters={sum(p.numel() for p in trainable):,}")
+        print(f"frozen_parameters={sum(p.numel() for p in model.parameters() if not p.requires_grad):,}")
+
+    model.to(device)
+    if world_size > 1:
+        model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
     optimizer = torch.optim.AdamW(
         trainable,
         lr=float(training["learning_rate"]),
@@ -139,24 +169,27 @@ def main():
         optimizer.load_state_dict(checkpoint["optimizer"])
     elif resume_checkpoint is not None:
         print("WARNING: resume checkpoint has no optimizer state; optimizer restarts fresh.")
+    del checkpoint
     epochs = 1 if args.smoke else (int(args.target_epoch) if args.target_epoch else int(training["epochs"]))
     if epochs < start_epoch:
         raise ValueError(f"target epoch {epochs} is before checkpoint epoch {start_epoch}")
     optimizer.zero_grad(set_to_none=True)
     for epoch in range(start_epoch, epochs):
+        if sampler is not None:
+            sampler.set_epoch(epoch)
         model.train()
         total_loss = 0.0
         for batch_index, (waveforms, labels) in enumerate(
             tqdm(
                 loader,
                 desc=f"epoch {epoch + 1}/{epochs}",
-                disable=not bool(training.get("progress_bar", True)),
+                disable=rank != 0 or not bool(training.get("progress_bar", True)),
             )
         ):
-            waveforms = waveforms.cuda(non_blocking=True)
-            labels = labels.cuda(non_blocking=True)
+            waveforms = waveforms.to(device, non_blocking=True)
+            labels = labels.to(device, non_blocking=True)
             timesteps = torch.randint(
-                0, diffusion.num_timesteps, (waveforms.shape[0],), device="cuda"
+                0, diffusion.num_timesteps, (waveforms.shape[0],), device=device
             )
             with torch.autocast("cuda", enabled=use_amp, dtype=torch.float16):
                 loss = diffusion.training_losses(
@@ -175,9 +208,14 @@ def main():
             if args.smoke:
                 break
         batches = 1 if args.smoke else len(loader)
-        print(f"epoch={epoch + 1} loss={total_loss / batches:.6f}")
+        mean_loss = torch.tensor(total_loss / batches, device=device)
+        if world_size > 1:
+            dist.all_reduce(mean_loss, op=dist.ReduceOp.SUM)
+            mean_loss /= world_size
+        if rank == 0:
+            print(f"epoch={epoch + 1} loss={mean_loss.item():.6f}")
         interval = int(config["output"].get("checkpoint_every_epochs", 10))
-        if not args.smoke and ((epoch + 1) % interval == 0 or epoch + 1 == epochs):
+        if rank == 0 and not args.smoke and ((epoch + 1) % interval == 0 or epoch + 1 == epochs):
             save_checkpoint(
                 checkpoint_dir / f"checkpoint_epoch_{epoch + 1:05d}.pt",
                 model,
@@ -187,7 +225,11 @@ def main():
                 config,
                 include_optimizer=bool(config["output"].get("save_optimizer_state", True)),
             )
+        if world_size > 1:
+            dist.barrier()
     dataset.close()
+    if world_size > 1:
+        dist.destroy_process_group()
 
 
 if __name__ == "__main__":

generate.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import argparse
+import os
 import sys
 from pathlib import Path
 
 import h5py
 import numpy as np
 import torch
+import torch.distributed as dist
 from tqdm import tqdm
 
 PROJECT_ROOT = Path(__file__).resolve().parent.parent
@@ -19,6 +21,19 @@ from diffusion import create_diffusion
 from models import EmptyConditionSWaG, SWaG
 
 
+def distributed_context() -> tuple[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:
+        dist.init_process_group("nccl")
+    return rank, local_rank, world_size
+
+
+def shard_path(output_path: Path, rank: int) -> Path:
+    return output_path.with_name(f".{output_path.name}.rank{rank}.partial.h5")
+
+
 
 def main() -> None:
     parser = argparse.ArgumentParser(description="Generate waveforms with SWaG")
@@ -40,52 +55,90 @@ def main() -> None:
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required for sampling")
 
+    rank, local_rank, world_size = distributed_context()
+    device = torch.device("cuda", local_rank)
+    torch.cuda.set_device(device)
+    if args.batch_size % world_size:
+        raise ValueError(f"batch-size must be divisible by world size ({world_size})")
+    local_batch = args.batch_size // world_size
     checkpoint_path = resolve_path(PROJECT_ROOT, str(args.checkpoint))
     output_path = resolve_path(PROJECT_ROOT, str(args.output))
     checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
     config = checkpoint["config"]
     model_class = EmptyConditionSWaG if int(config["model"].get("condition_slot_count", 0)) == 8 else SWaG
-    model = model_class(**config["model"]).cuda().eval()
+    model = model_class(**config["model"]).to(device).eval()
     model.load_state_dict(checkpoint["ema"], strict=True)
+    del checkpoint
     diffusion_config = dict(config["diffusion"])
     respacing = f"ddim{args.sampling_steps}" if args.sampler == "ddim" else str(args.sampling_steps)
     diffusion = create_diffusion(timestep_respacing=respacing, **diffusion_config)
-    generator = torch.Generator(device="cuda").manual_seed(args.seed)
-    output_path.parent.mkdir(parents=True, exist_ok=True)
+    generator = torch.Generator(device=device).manual_seed(args.seed + rank)
+    if rank == 0:
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        for shard_rank in range(world_size):
+            shard_path(output_path, shard_rank).unlink(missing_ok=True)
+    if world_size > 1:
+        dist.barrier()
     shape = (args.num_samples, int(config["model"]["length"]), int(config["model"]["in_channels"]))
-    with h5py.File(output_path, "w") as handle:
-        output = handle.create_dataset("data", shape=shape, dtype="float32", compression="gzip", shuffle=True)
-        output.attrs["checkpoint"] = str(checkpoint_path)
-        output.attrs["sampling_steps"] = args.sampling_steps
-        output.attrs["sampler"] = args.sampler
-        output.attrs["seed"] = args.seed
-        output.attrs["clip_denoised"] = args.clip_denoised
-        for start in tqdm(range(0, args.num_samples, args.batch_size), desc="sampling"):
-            count = min(args.batch_size, args.num_samples - start)
+    assigned = np.array_split(np.arange(args.num_samples, dtype=np.int64), world_size)[rank]
+    local_shape = (len(assigned), shape[1], shape[2])
+    with h5py.File(shard_path(output_path, rank), "w") as handle:
+        output = handle.create_dataset("data", shape=local_shape, dtype="float32", compression="gzip", shuffle=True)
+        handle.create_dataset("indices", data=assigned, dtype="int64")
+        for start in tqdm(
+            range(0, len(assigned), local_batch),
+            desc=f"sampling rank {rank}",
+            disable=rank != 0,
+        ):
+            count = min(local_batch, len(assigned) - start)
             noise = torch.randn(
                 count,
                 int(config["model"]["in_channels"]),
                 int(config["model"]["length"]),
-                device="cuda",
+                device=device,
                 generator=generator,
             )
             sample_loop = diffusion.ddim_sample_loop if args.sampler == "ddim" else diffusion.p_sample_loop
             model_kwargs = None
             if int(config["model"].get("condition_slot_count", 0)) == 8:
                 model_kwargs = {
-                    "conditions": torch.zeros(count, 8, device="cuda", dtype=torch.float32)
+                    "conditions": torch.zeros(count, 8, device=device, dtype=torch.float32)
                 }
             samples = sample_loop(
                 model,
                 noise.shape,
                 noise=noise,
-                device=torch.device("cuda"),
+                device=device,
                 progress=False,
                 clip_denoised=args.clip_denoised,
                 model_kwargs=model_kwargs,
             )
             output[start : start + count] = samples.transpose(1, 2).cpu().numpy().astype(np.float32)
 
+    if world_size > 1:
+        dist.barrier()
+    if rank == 0:
+        partial = output_path.with_name(f".{output_path.name}.partial")
+        partial.unlink(missing_ok=True)
+        with h5py.File(partial, "w") as handle:
+            output = handle.create_dataset("data", shape=shape, dtype="float32", compression="gzip", shuffle=True)
+            handle.attrs["checkpoint"] = str(checkpoint_path)
+            handle.attrs["sampling_steps"] = args.sampling_steps
+            handle.attrs["sampler"] = args.sampler
+            handle.attrs["seed"] = args.seed
+            handle.attrs["clip_denoised"] = args.clip_denoised
+            handle.attrs["world_size"] = world_size
+            for shard_rank in range(world_size):
+                current = shard_path(output_path, shard_rank)
+                with h5py.File(current, "r") as shard:
+                    indices = np.asarray(shard["indices"], dtype=np.int64)
+                    output[indices] = shard["data"][:]
+                current.unlink()
+        partial.replace(output_path)
+    if world_size > 1:
+        dist.barrier()
+        dist.destroy_process_group()
+
 
 if __name__ == "__main__":
     main()

run_all_local.sh

@@ -4,5 +4,4 @@ set -euo pipefail
 SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
 
 bash "${SCRIPT_DIR}/run_download_prepare.sh"
-bash "${SCRIPT_DIR}/run_train_local.sh"
-
+bash "${SCRIPT_DIR}/run_train_local.sh" "$@"

run_generate_100.sh

@@ -7,14 +7,38 @@ PYTHON_BIN="${PYTHON_BIN:-python}"
 CHECKPOINT="${CHECKPOINT:-${SCRIPT_DIR}/results/stead_empty8_local/checkpoints/checkpoint_epoch_00015.pt}"
 RESULT_DIR="${RESULT_DIR:-${SCRIPT_DIR}/results/stead_empty8_local}"
 OUTPUT="${OUTPUT:-${RESULT_DIR}/generated_100/generated_ddpm25_seed2026.h5}"
+NUM_GPUS="${NUM_GPUS:-8}"
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --num-gpus)
+      [[ $# -ge 2 ]] || { echo "--num-gpus requires a value" >&2; exit 2; }
+      NUM_GPUS="$2"
+      shift 2
+      ;;
+    --num-gpus=*) NUM_GPUS="${1#*=}"; shift ;;
+    *) echo "Usage: bash Train/run_generate_100.sh [--num-gpus N]" >&2; exit 2 ;;
+  esac
+done
+[[ "${NUM_GPUS}" =~ ^[1-9][0-9]*$ ]] || { echo "NUM_GPUS must be a positive integer" >&2; exit 2; }
 
 cd "${PROJECT_ROOT}"
 [[ -f "${CHECKPOINT}" ]] || { echo "Missing checkpoint: ${CHECKPOINT}" >&2; exit 1; }
-exec "${PYTHON_BIN}" Train/generate.py \
-  --checkpoint "${CHECKPOINT}" \
-  --output "${OUTPUT}" \
-  --num-samples 100 \
-  --batch-size 256 \
-  --sampling-steps 25 \
-  --sampler ddpm \
+ARGS=(
+  Train/generate.py
+  --checkpoint "${CHECKPOINT}"
+  --output "${OUTPUT}"
+  --num-samples 100
+  --batch-size 256
+  --sampling-steps 25
+  --sampler ddpm
   --seed 2026
+)
+echo "Generating with ${NUM_GPUS} GPU(s); global generation batch size is 256"
+if [[ "${NUM_GPUS}" -eq 1 ]]; then
+  exec "${PYTHON_BIN}" "${ARGS[@]}"
+fi
+exec "${PYTHON_BIN}" -m torch.distributed.run \
+  --standalone \
+  --nproc_per_node="${NUM_GPUS}" \
+  "${ARGS[@]}"

run_train_local.sh

@@ -6,6 +6,27 @@ PROJECT_ROOT="$(cd -- "${SCRIPT_DIR}/.." && pwd)"
 PYTHON_BIN="${PYTHON_BIN:-python}"
 CONFIG="${CONFIG:-${SCRIPT_DIR}/configs/stead_empty8_local.yaml}"
 TRAIN_H5="${PROJECT_ROOT}/Train/data/processed/from_raw/stead/train/stead_100hz_60s_train.h5"
+NUM_GPUS="${NUM_GPUS:-8}"
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --num-gpus)
+      [[ $# -ge 2 ]] || { echo "--num-gpus requires a value" >&2; exit 2; }
+      NUM_GPUS="$2"
+      shift 2
+      ;;
+    --num-gpus=*)
+      NUM_GPUS="${1#*=}"
+      shift
+      ;;
+    *)
+      echo "Unknown argument: $1" >&2
+      echo "Usage: bash Train/run_train_local.sh [--num-gpus N]" >&2
+      exit 2
+      ;;
+  esac
+done
+[[ "${NUM_GPUS}" =~ ^[1-9][0-9]*$ ]] || { echo "NUM_GPUS must be a positive integer" >&2; exit 2; }
 
 cd "${PROJECT_ROOT}"
 [[ -f "${TRAIN_H5}" ]] || {
@@ -15,5 +36,11 @@ cd "${PROJECT_ROOT}"
 }
 
 export PYTHONUNBUFFERED=1
-exec "${PYTHON_BIN}" Train/training/train.py --config "${CONFIG}"
-
+echo "Training with ${NUM_GPUS} GPU(s); global batch size is read from ${CONFIG}"
+if [[ "${NUM_GPUS}" -eq 1 ]]; then
+  exec "${PYTHON_BIN}" Train/training/train.py --config "${CONFIG}"
+fi
+exec "${PYTHON_BIN}" -m torch.distributed.run \
+  --standalone \
+  --nproc_per_node="${NUM_GPUS}" \
+  Train/training/train.py --config "${CONFIG}"

training/train.py

@@ -195,8 +195,12 @@ def main() -> None:
                 step += 1
             if rank == 0:
                 progress.set_postfix(loss=f"{running_loss / (batch_index + 1):.5f}")
+        mean_loss = torch.tensor(running_loss / max(len(loader), 1), device=device)
+        if world_size > 1:
+            dist.all_reduce(mean_loss, op=dist.ReduceOp.SUM)
+            mean_loss /= world_size
         if rank == 0:
-            metrics = {"epoch": epoch + 1, "step": step, "loss": running_loss / max(len(loader), 1)}
+            metrics = {"epoch": epoch + 1, "step": step, "loss": mean_loss.item()}
             with (output_dir / "metrics.jsonl").open("a", encoding="utf-8") as handle:
                 handle.write(json.dumps(metrics) + "\n")
             if (epoch + 1) % checkpoint_interval == 0 or epoch + 1 == epochs:
@@ -204,6 +208,8 @@ def main() -> None:
                     checkpoint_dir / f"checkpoint_epoch_{epoch + 1:05d}.pt",
                     model, ema, optimizer, epoch + 1, step, config,
                 )
+        if world_size > 1:
+            dist.barrier()
 
     dataset.close()
     if world_size > 1: