Merge branch 'main' of github.com:pesser/stable-diffusion into main

configs/stable-diffusion/v1_laionhr.yaml

@@ -0,0 +1,135 @@
+model:
+  base_learning_rate: 1.0e-04
+  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    image_size: 64
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: crossattn
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+
+    scheduler_config: # 10000 warmup steps
+      target: ldm.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 1 ] # NOTE for resuming. use 10000 if starting from scratch
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    unet_config:
+      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
+
+
+data:
+  target: ldm.data.laion.WebDataModuleFromConfig
+  params:
+    tar_base: "pipe:aws s3 cp s3://s-datasets/laion-high-resolution/"
+    batch_size: 4
+    num_workers: 4
+    multinode: True
+    train:
+      shards: '{00000..17279}.tar -'
+      shuffle: 10000
+      image_key: jpg
+      image_transforms:
+      - target: torchvision.transforms.Resize
+        params:
+          size: 512
+          interpolation: 3
+      - target: torchvision.transforms.RandomCrop
+        params:
+          size: 512
+
+    # NOTE use enough shards to avoid empty validation loops in workers
+    validation:
+      shards: '{17280..17535}.tar -'
+      shuffle: 0
+      image_key: jpg
+      image_transforms:
+      - target: torchvision.transforms.Resize
+        params:
+          size: 512
+          interpolation: 3
+      - target: torchvision.transforms.CenterCrop
+        params:
+          size: 512
+
+
+lightning:
+  find_unused_parameters: False
+
+  modelcheckpoint:
+    params:
+      every_n_train_steps: 5000
+
+  callbacks:
+    image_logger:
+      target: main.ImageLogger
+      params:
+        batch_frequency: 5000
+        max_images: 4
+        increase_log_steps: False
+        log_first_step: False
+        log_images_kwargs:
+          use_ema_scope: False
+          inpaint: False
+          plot_progressive_rows: False
+          plot_diffusion_rows: False
+          N: 4
+          unconditional_guidance_scale: 3.0
+          unconditional_guidance_label: [""]
+
+  trainer:
+    benchmark: True
+    val_check_interval: 5000000  # really sorry
+    num_sanity_val_steps: 0
+    accumulate_grad_batches: 2

ldm/data/laion.py

@@ -105,6 +105,7 @@ def dict_collation_fn(samples, combine_tensors=True, combine_scalars=True):
 class WebDataModuleFromConfig(pl.LightningDataModule):
     def __init__(self, tar_base, batch_size, train=None, validation=None,
                  test=None, num_workers=4, multinode=True, min_size=None,
+                 max_pwatermark=1.0,
                  **kwargs):
         super().__init__(self)
         print(f'Setting tar base to {tar_base}')
@@ -116,6 +117,7 @@ class WebDataModuleFromConfig(pl.LightningDataModule):
         self.test = test
         self.multinode = multinode
         self.min_size = min_size  # filter out very small images
+        self.max_pwatermark = max_pwatermark # filter out watermarked images
 
     def make_loader(self, dataset_config, train=True):
         if 'image_transforms' in dataset_config:
@@ -184,7 +186,7 @@ class WebDataModuleFromConfig(pl.LightningDataModule):
         if self.min_size is None:
             return True
         try:
-            return x['json']['original_width'] >= self.min_size and x['json']['original_height'] >= self.min_size
+            return x['json']['original_width'] >= self.min_size and x['json']['original_height'] >= self.min_size and x['json']['pwatermark'] <= self.max_pwatermark
         except Exception:
             return False
 
@@ -300,8 +302,7 @@ def example01():
         print("next epoch.")
 
 
-if __name__ == "__main__":
+def example02():
-    #example01()
     from omegaconf import OmegaConf
     from torch.utils.data.distributed import DistributedSampler
     from torch.utils.data import IterableDataset
@@ -318,3 +319,82 @@ if __name__ == "__main__":
         print(batch.keys())
         print(batch["jpg"].shape)
         break
+
+
+def example03():
+    # improved aesthetics
+    tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{000000..060207}.tar -"
+    dataset = wds.WebDataset(tars)
+
+    def filter_keys(x):
+        try:
+            return ("jpg" in x) and ("txt" in x)
+        except Exception:
+            return False
+
+    def filter_size(x):
+        try:
+            return x['json']['original_width'] >= 512 and x['json']['original_height'] >= 512
+        except Exception:
+            return False
+
+    def filter_watermark(x):
+        try:
+            return x['json']['pwatermark'] < 0.5
+        except Exception:
+            return False
+
+    dataset = (dataset
+                .select(filter_keys)
+                .decode('pil', handler=wds.warn_and_continue))
+    n_total = 0
+    n_large = 0
+    n_large_nowm = 0
+    for i, example in enumerate(dataset):
+        n_total += 1
+        if filter_size(example):
+            n_large += 1
+            if filter_watermark(example):
+                n_large_nowm += 1
+
+        if i%500 == 0:
+            print(i)
+            print(f"Large: {n_large}/{n_total} | {n_large/n_total*100:.2f}%")
+            if n_large > 0:
+                print(f"No Watermark: {n_large_nowm}/{n_large} | {n_large_nowm/n_large*100:.2f}%")
+
+
+
+def example04():
+    # improved aesthetics
+    for i_shard in range(60208)[::-1]:
+        print(i_shard)
+        tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{:06}.tar -".format(i_shard)
+        dataset = wds.WebDataset(tars)
+
+        def filter_keys(x):
+            try:
+                return ("jpg" in x) and ("txt" in x)
+            except Exception:
+                return False
+
+        def filter_size(x):
+            try:
+                return x['json']['original_width'] >= 512 and x['json']['original_height'] >= 512
+            except Exception:
+                return False
+
+        dataset = (dataset
+                    .select(filter_keys)
+                    .decode('pil', handler=wds.warn_and_continue))
+        try:
+            example = next(iter(dataset))
+        except Exception:
+            print(f"Error @ {i_shard}")
+
+
+if __name__ == "__main__":
+    #example01()
+    #example02()
+    example03()
+    #example04()

scripts/cmd_on_new_ckpt.py

@@ -0,0 +1,32 @@
+import os
+import subprocess
+import time
+import fire
+
+
+class Checker(object):
+    def __init__(self, filename, interval=60):
+        self._cached_stamp = 0
+        self.filename = filename
+        self.interval = interval
+    
+    def check(self, cmd):
+        while True:
+            stamp = os.stat(self.filename).st_mtime
+            if stamp != self._cached_stamp:
+                self._cached_stamp = stamp
+                print(f"{self.__class__.__name__}: Detected a new file at {self.filename}, running evaluation commands on it.")
+                subprocess.run(cmd, shell=True)
+            else:
+                time.sleep(self.interval)
+
+
+def run(filename, cmd):
+    checker = Checker(filename, interval=60)
+    checker.check(cmd)
+
+
+if __name__ == "__main__":
+    fire.Fire(run)
+                
+

scripts/printckpt.py

@@ -0,0 +1,16 @@
+import os
+import torch
+import fire
+
+
+def printit(p):
+    print(f"printin' in path: {p}")
+    size_initial = os.path.getsize(p)
+    nsd = dict()
+    sd = torch.load(p, map_location="cpu")
+    if "global_step" in sd:
+        print(f"This is global step {sd['global_step']}.")
+
+
+if __name__ == "__main__":
+    fire.Fire(printit)

scripts/prune-ckpt.py

@@ -14,6 +14,8 @@ def prune_it(p):
             nsd[k] = sd[k]
     else:
         print(f"removing optimizer states for path {p}")
+    if "global_step" in sd:
+        print(f"This is global step {sd['global_step']}.")
     fn = f"{os.path.splitext(p)[0]}-pruned.ckpt"
     print(f"saving pruned checkpoint at: {fn}")
     torch.save(nsd, fn)

scripts/slurm/v1_iahr_torch111/launcher.sh

@@ -0,0 +1,46 @@
+#!/bin/bash
+
+# mpi version for node rank
+H=`hostname`
+THEID=`echo -e $HOSTNAMES  | python3 -c "import sys;[sys.stdout.write(str(i)) for i,line in enumerate(next(sys.stdin).split(' ')) if line.strip() == '$H'.strip()]"`
+export NODE_RANK=${THEID}
+echo THEID=$THEID
+
+echo "##########################################"
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo NODE_RANK=${NODE_RANK}
+echo WORLD_SIZE=${WORLD_SIZE}
+echo "##########################################"
+# debug environment worked great so we stick with it
+# no magic there, just a miniconda python=3.9, pytorch=1.12, cudatoolkit=11.3
+# env with pip dependencies from stable diffusion's requirements.txt
+eval "$(/fsx/stable-diffusion/debug/miniconda3/bin/conda shell.bash hook)"
+#conda activate stable
+conda activate torch111
+cd /fsx/stable-diffusion/stable-diffusion
+
+CONFIG="/fsx/stable-diffusion/stable-diffusion/configs/stable-diffusion/v1_improvedaesthetics.yaml"
+
+# resume and set new seed to reshuffle data
+#EXTRA="--seed 718 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/checkpoints2/v1pp/v1pp-flatline.ckpt"
+EXTRA="--seed 718 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-22T07-45-07_v1_improvedaesthetics/checkpoints/last.ckpt"
+
+# only images >= 512 and pwatermark <= 0.4999
+EXTRA="${EXTRA} data.params.min_size=512 data.params.max_pwatermark=0.4999"
+
+# postfix
+EXTRA="${EXTRA} -f v1_iahr_torch111"
+
+# time to decay
+#EXTRA="${EXTRA} model.params.scheduler_config.params.cycle_lengths=[50000] model.params.scheduler_config.params.f_min=[1e-6]"
+
+# custom logdir
+#EXTRA="${EXTRA} --logdir rlogs"
+
+# debugging
+#EXTRA="${EXTRA} -d True lightning.callbacks.image_logger.params.batch_frequency=50"
+
+/bin/bash /fsx/stable-diffusion/stable-diffusion/scripts/test_gpu.sh
+
+python main.py --base $CONFIG --gpus 0,1,2,3,4,5,6,7 -t --num_nodes ${WORLD_SIZE} --scale_lr False $EXTRA

scripts/slurm/v1_iahr_torch111/sbatch.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+#SBATCH --partition=compute-od-gpu
+#SBATCH --job-name=stable-diffusion-v1-iahr-torch111
+#SBATCH --nodes 32
+#SBATCH --ntasks-per-node 1
+#SBATCH --cpus-per-gpu=4
+#SBATCH --gres=gpu:8
+#SBATCH --exclusive
+#SBATCH --output=%x_%j.out
+#SBATCH --comment "Key=Monitoring,Value=ON"
+
+module load intelmpi
+source /opt/intel/mpi/latest/env/vars.sh
+export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+#export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-inst
+#all/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+export NCCL_PROTO=simple
+export PATH=/opt/amazon/efa/bin:$PATH
+export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
+export NCCL_DEBUG=info
+export PYTHONFAULTHANDLER=1
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+
+# sent to sub script
+export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=12802
+export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
+export WORLD_SIZE=$COUNT_NODE
+
+echo go $COUNT_NODE
+echo $HOSTNAMES
+echo $WORLD_SIZE
+
+mpirun -n $COUNT_NODE -perhost 1 /fsx/stable-diffusion/stable-diffusion/scripts/slurm/v1_iahr_torch111/launcher.sh

scripts/slurm/v1_improvedaesthetics/launcher.sh

@@ -22,7 +22,10 @@ cd /fsx/stable-diffusion/stable-diffusion
 CONFIG="/fsx/stable-diffusion/stable-diffusion/configs/stable-diffusion/v1_improvedaesthetics.yaml"
 
 # resume and set new seed to reshuffle data
-EXTRA="--seed 714 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/logs/2022-07-11T20-16-11_txt2img-1p4B-multinode-clip-encoder-high-res-512_improvedaesthetic/checkpoints/last.ckpt"
+#EXTRA="--seed 714 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/logs/2022-07-11T20-16-11_txt2img-1p4B-multinode-clip-encoder-high-res-512_improvedaesthetic/checkpoints/last.ckpt"
+EXTRA="--seed 715 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/logs/2022-07-14T23-26-13_v1_improvedaesthetics/checkpoints/last.ckpt"
+# time to decay
+EXTRA="${EXTRA} model.params.scheduler_config.params.cycle_lengths=[50000] model.params.scheduler_config.params.f_min=[1e-6]"
 
 # custom logdir
 #EXTRA="${EXTRA} --logdir rlogs"

scripts/slurm/v1_improvedaesthetics_torch111/launcher.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+
+# mpi version for node rank
+H=`hostname`
+THEID=`echo -e $HOSTNAMES  | python3 -c "import sys;[sys.stdout.write(str(i)) for i,line in enumerate(next(sys.stdin).split(' ')) if line.strip() == '$H'.strip()]"`
+export NODE_RANK=${THEID}
+echo THEID=$THEID
+
+echo "##########################################"
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo NODE_RANK=${NODE_RANK}
+echo WORLD_SIZE=${WORLD_SIZE}
+echo "##########################################"
+# debug environment worked great so we stick with it
+# no magic there, just a miniconda python=3.9, pytorch=1.12, cudatoolkit=11.3
+# env with pip dependencies from stable diffusion's requirements.txt
+eval "$(/fsx/stable-diffusion/debug/miniconda3/bin/conda shell.bash hook)"
+conda activate torch111
+cd /fsx/stable-diffusion/stable-diffusion
+
+CONFIG="/fsx/stable-diffusion/stable-diffusion/configs/stable-diffusion/v1_improvedaesthetics.yaml"
+
+# resume and set new seed to reshuffle data
+#EXTRA="--seed 714 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/logs/2022-07-11T20-16-11_txt2img-1p4B-multinode-clip-encoder-high-res-512_improvedaesthetic/checkpoints/last.ckpt"
+#EXTRA="--seed 715 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/logs/2022-07-14T23-26-13_v1_improvedaesthetics/checkpoints/last.ckpt"
+#EXTRA="--seed 716 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-18T23-11-11_v1_improvedaesthetics/checkpoints/last.ckpt"
+EXTRA="--seed 717 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-19T19-03-04_v1_improvedaesthetics/checkpoints/last.ckpt"
+# time to decay
+EXTRA="${EXTRA} model.params.scheduler_config.params.cycle_lengths=[50000] model.params.scheduler_config.params.f_min=[1e-6]"
+
+# custom logdir
+#EXTRA="${EXTRA} --logdir rlogs"
+
+# debugging
+#EXTRA="${EXTRA} -d True lightning.callbacks.image_logger.params.batch_frequency=50"
+
+python main.py --base $CONFIG --gpus 0,1,2,3,4,5,6,7 -t --num_nodes ${WORLD_SIZE} --scale_lr False $EXTRA

scripts/slurm/v1_improvedaesthetics_torch111/sbatch.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+#SBATCH --partition=compute-od-gpu
+#SBATCH --job-name=stable-diffusion-v1-improvedaesthetics-torch111
+#SBATCH --nodes 20
+#SBATCH --ntasks-per-node 1
+#SBATCH --cpus-per-gpu=4
+#SBATCH --gres=gpu:8
+#SBATCH --exclusive
+#SBATCH --output=%x_%j.out
+#SBATCH --comment "Key=Monitoring,Value=ON"
+
+module load intelmpi
+source /opt/intel/mpi/latest/env/vars.sh
+export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-inst
+all/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+export NCCL_PROTO=simple
+export PATH=/opt/amazon/efa/bin:$PATH
+export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
+export NCCL_DEBUG=info
+export PYTHONFAULTHANDLER=1
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+
+# sent to sub script
+export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=12802
+export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
+export WORLD_SIZE=$COUNT_NODE
+
+echo go $COUNT_NODE
+echo $HOSTNAMES
+echo $WORLD_SIZE
+
+mpirun -n $COUNT_NODE -perhost 1 /fsx/stable-diffusion/stable-diffusion/scripts/slurm/v1_improvedaesthetics_torch111/launcher.sh

scripts/slurm/v1_laionhr_torch111/launcher.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+
+# mpi version for node rank
+H=`hostname`
+THEID=`echo -e $HOSTNAMES  | python3 -c "import sys;[sys.stdout.write(str(i)) for i,line in enumerate(next(sys.stdin).split(' ')) if line.strip() == '$H'.strip()]"`
+export NODE_RANK=${THEID}
+echo THEID=$THEID
+
+echo "##########################################"
+echo MASTER_ADDR=${MASTER_ADDR}
+echo MASTER_PORT=${MASTER_PORT}
+echo NODE_RANK=${NODE_RANK}
+echo WORLD_SIZE=${WORLD_SIZE}
+echo "##########################################"
+# debug environment worked great so we stick with it
+# no magic there, just a miniconda python=3.9, pytorch=1.12, cudatoolkit=11.3
+# env with pip dependencies from stable diffusion's requirements.txt
+eval "$(/fsx/stable-diffusion/debug/miniconda3/bin/conda shell.bash hook)"
+conda activate torch111
+cd /fsx/stable-diffusion/stable-diffusion
+
+CONFIG="/fsx/stable-diffusion/stable-diffusion/configs/stable-diffusion/v1_laionhr.yaml"
+
+# resume and set new seed to reshuffle data
+EXTRA="--seed 718 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/checkpoints2/v1pp/v1pp-flatline.ckpt"
+
+# time to decay
+#EXTRA="${EXTRA} model.params.scheduler_config.params.cycle_lengths=[50000] model.params.scheduler_config.params.f_min=[1e-6]"
+
+# custom logdir
+#EXTRA="${EXTRA} --logdir rlogs"
+
+# debugging
+#EXTRA="${EXTRA} -d True lightning.callbacks.image_logger.params.batch_frequency=50"
+
+python main.py --base $CONFIG --gpus 0,1,2,3,4,5,6,7 -t --num_nodes ${WORLD_SIZE} --scale_lr False $EXTRA

scripts/slurm/v1_laionhr_torch111/sbatch.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+#SBATCH --partition=compute-od-gpu
+#SBATCH --job-name=stable-diffusion-v1-laionhr-torch111
+#SBATCH --nodes 20
+#SBATCH --ntasks-per-node 1
+#SBATCH --cpus-per-gpu=4
+#SBATCH --gres=gpu:8
+#SBATCH --exclusive
+#SBATCH --output=%x_%j.out
+#SBATCH --comment "Key=Monitoring,Value=ON"
+
+module load intelmpi
+source /opt/intel/mpi/latest/env/vars.sh
+export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-inst
+all/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
+export NCCL_PROTO=simple
+export PATH=/opt/amazon/efa/bin:$PATH
+export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"
+export FI_EFA_FORK_SAFE=1
+export FI_LOG_LEVEL=1
+export FI_EFA_USE_DEVICE_RDMA=1 # use for p4dn
+export NCCL_DEBUG=info
+export PYTHONFAULTHANDLER=1
+export CUDA_LAUNCH_BLOCKING=0
+export OMPI_MCA_mtl_base_verbose=1
+export FI_EFA_ENABLE_SHM_TRANSFER=0
+export FI_PROVIDER=efa
+export FI_EFA_TX_MIN_CREDITS=64
+export NCCL_TREE_THRESHOLD=0
+
+# sent to sub script
+export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=12802
+export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
+export WORLD_SIZE=$COUNT_NODE
+
+echo go $COUNT_NODE
+echo $HOSTNAMES
+echo $WORLD_SIZE
+
+mpirun -n $COUNT_NODE -perhost 1 /fsx/stable-diffusion/stable-diffusion/scripts/slurm/v1_laionhr_torch111/launcher.sh

scripts/slurm/v3_pretraining/launcher.sh

@@ -16,13 +16,16 @@ echo "##########################################"
 # no magic there, just a miniconda python=3.9, pytorch=1.12, cudatoolkit=11.3
 # env with pip dependencies from stable diffusion's requirements.txt
 eval "$(/fsx/stable-diffusion/debug/miniconda3/bin/conda shell.bash hook)"
-conda activate stable
+#conda activate stable
+# torch 1.11 to avoid bug in ckpt restoring
+conda activate torch111
 cd /fsx/stable-diffusion/stable-diffusion
 
 CONFIG=configs/stable-diffusion/v3_pretraining.yaml
 
 # resume and set new seed to reshuffle data
-EXTRA="--seed 714 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/rlogs/2022-07-11T22-57-10_txt2img-v2-clip-encoder-improved_aesthetics-256/checkpoints/last.ckpt"
+#EXTRA="--seed 714 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/rlogs/2022-07-11T22-57-10_txt2img-v2-clip-encoder-improved_aesthetics-256/checkpoints/last.ckpt"
+EXTRA="--seed 715 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-14T21-03-49_txt2img-v2-clip-encoder-improved_aesthetics-256/checkpoints/last.ckpt"
 
 # custom logdir
 #EXTRA="${EXTRA} --logdir rlogs"
@@ -30,4 +33,7 @@ EXTRA="--seed 714 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/
 # debugging
 #EXTRA="${EXTRA} -d True lightning.callbacks.image_logger.params.batch_frequency=50"
 
+# detect bad gpus early on
+/bin/bash /fsx/stable-diffusion/stable-diffusion/scripts/test_gpu.sh
+
 python main.py --base $CONFIG --gpus 0,1,2,3,4,5,6,7 -t --num_nodes ${WORLD_SIZE} --scale_lr False $EXTRA

scripts/slurm/v3_pretraining/sbatch.sh

@@ -11,8 +11,7 @@
 
 module load intelmpi
 source /opt/intel/mpi/latest/env/vars.sh
-export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-inst
+export LD_LIBRARY_PATH=/opt/aws-ofi-nccl/lib:/opt/amazon/efa/lib64:/usr/local/cuda-11.0/efa/lib:/usr/local/cuda-11.0/lib:/usr/local/cuda-11.0/lib64:/usr/local/cuda-11.0:/opt/nccl/build/lib:/opt/aws-ofi-nccl-install/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
-all/lib:/opt/aws-ofi-nccl/lib:$LD_LIBRARY_PATH
 export NCCL_PROTO=simple
 export PATH=/opt/amazon/efa/bin:$PATH
 export LD_PRELOAD="/opt/nccl/build/lib/libnccl.so"

scripts/test_gpu.py

@@ -0,0 +1,30 @@
+import socket
+try:
+    import torch
+    n_gpus = torch.cuda.device_count()
+    print(f"checking {n_gpus} gpus.")
+    for i_gpu in range(n_gpus):
+        print(i_gpu)
+        device = torch.device(f"cuda:{i_gpu}")
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cudnn.deterministic = False
+        torch.backends.cudnn.allow_tf32 = True
+        data = torch.randn([4, 640, 32, 32], dtype=torch.float, device=device, requires_grad=True)
+        net = torch.nn.Conv2d(640, 640, kernel_size=[3, 3], padding=[1, 1], stride=[1, 1], dilation=[1, 1], groups=1)
+        net = net.to(device=device).float()
+        out = net(data)
+        out.backward(torch.randn_like(out))
+        torch.cuda.synchronize()
+except RuntimeError as err:
+    import requests
+    import datetime
+    import os
+    device = os.environ.get("CUDA_VISIBLE_DEVICES", "?")
+    hostname = socket.gethostname()
+    ts = datetime.datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S')
+    resp = requests.get('http://169.254.169.254/latest/meta-data/instance-id')
+    print(f'ERROR at {ts} on {hostname}/{resp.text} (CUDA_VISIBLE_DEVICES={device}): {type(err).__name__}: {err}', flush=True)
+    raise err
+else:
+    print(f"checked {socket.gethostname()}")

scripts/test_gpu.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+eval "$(/fsx/stable-diffusion/debug/miniconda3/bin/conda shell.bash hook)"
+conda activate stable
+cd /fsx/stable-diffusion/stable-diffusion
+python scripts/test_gpu.py