SD-2-base-512.yaml

project: # Insert wandb project name
name: # Insert wandb run name
seed: 17
eval_first: false
algorithms:
  ema:
    _target_: diffusion.algorithms.ema.EMA
    half_life: null
    smoothing: 0.9999
    update_interval: 1ba
    ema_start: 800000ba
  low_precision_groupnorm:
    attribute: unet
    precision: amp_fp16
  low_precision_layernorm:
    attribute: unet
    precision: amp_fp16
model:
  _target_: diffusion.models.models.stable_diffusion_2
  pretrained: false
  precomputed_latents: true
  encode_latents_in_fp16: true
  fsdp: true
  val_metrics:
    - _target_: torchmetrics.MeanSquaredError
dataset:
  train_batch_size: 2048 # Global training batch size
  eval_batch_size: 1024  # Global evaluation batch size
  train_dataset:
    _target_: diffusion.datasets.laion.laion.build_streaming_laion_dataloader
    remote:
      # Path to object store bucket(s)
    local:
      # Path to corresponding local dataset(s)
    tokenizer_name_or_path: stabilityai/stable-diffusion-2-base
    caption_drop_prob: 0.1
    resize_size: 512
    drop_last: true
    shuffle: true
    prefetch_factor: 2
    num_workers: 8
    persistent_workers: true
    pin_memory: true
    download_timeout: 300
    num_canonical_nodes: 64
  eval_dataset:
    _target_: diffusion.datasets.coco.coco_captions.build_streaming_cocoval_dataloader
    remote: # Path to object store bucket
    local: # Path to local dataset cache
    resize_size: 512
    prefetch_factor: 2
    num_workers: 8
    persistent_workers: True
    pin_memory: True
optimizer:
  _target_: torch.optim.AdamW
  lr: 1.0e-4
  weight_decay: 0.01
scheduler:
  _target_: composer.optim.LinearWithWarmupScheduler
  t_warmup: 10000ba
  alpha_f: 1.0
logger:
  wandb:
    _target_: composer.loggers.wandb_logger.WandBLogger
    name: ${name}
    project: ${project}
    group: ${name}
callbacks:
  speed_monitor:
    _target_: composer.callbacks.speed_monitor.SpeedMonitor
    window_size: 10
  lr_monitor:
    _target_: composer.callbacks.lr_monitor.LRMonitor
  memory_monitor:
    _target_: composer.callbacks.memory_monitor.MemoryMonitor
  runtime_estimator:
    _target_: composer.callbacks.runtime_estimator.RuntimeEstimator
  optimizer_monitor:
    _target_: composer.callbacks.OptimizerMonitor
trainer:
  _target_: composer.Trainer
  device: gpu
  max_duration: 850000ba
  eval_interval: 10000ba
  device_train_microbatch_size: 16
  run_name: ${name}
  seed: ${seed}
  save_folder:  # Insert path to save folder or bucket
  save_interval: 10000ba
  save_overwrite: true
  autoresume: false
  fsdp_config:
    sharding_strategy: "SHARD_GRAD_OP"