Textual Inversion Training on M1 (works!)

[tmm1]

WIP HERE: development...tmm1:dev-train-m1

---

I started experimenting with running main.py on M1 and wanted to document some immediate issues.

Looks like we need a newer pytorch-lightning for MPS. Currently using 1.6.5 but latest is 1.7.5

However bumping it causes this error:

AttributeError: module 'pytorch_lightning.loggers' has no a
ttribute 'TestTubeLogger'. Did you mean: 'NeptuneLogger'?

which is because TestTubeLogger was deprecated: Lightning-AI/pytorch-lightning#13958 (comment)

[lstein]

I started working with the training functionality last night as well and ran into problems on CUDA. The textual inversion modifications to ddpm.py seem to have adversely affected vanilla training and we'll have to do a careful comparison with the original CompViz implementation in order to isolate the conflicts.

@tmm1, have you tried main.py on M1 using any of the other (multitudinous) forks? If so, any success?

[tmm1]

If there's a fork that advertises M1 training support I would be happy to try it. I have not seen one, but I have not looked much either. My understanding was that most of the M1 work was happening here.

[lstein]

@Any-Winter-4079, when you've finished the latest round of code tweaking, could you have a look at training? It seems to be messed up on M1.

[tmm1]

I made some progress today, and was able to get through all the setup and start training to send commands to the mps backend: development...tmm1:dev-train-m1

Currently stuck here:

...
Global seed set to 23
Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/1
----------------------------------------------------------------------------------------------------
distributed_backend=gloo
All distributed processes registered. Starting with 1 processes
----------------------------------------------------------------------------------------------------

Summoning checkpoint.
Traceback (most recent call last):
  File "/Users/tmm1/code/stable-diffusion/./main.py", line 946, in <module>
    trainer.fit(model, data)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 696, in fit
    self._call_and_handle_interrupt(
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 648, in _call_and_handle_interrupt
    return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, **kwargs)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 93, in launch
    return function(*args, **kwargs)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 735, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1147, in _run
    self.strategy.setup(self)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 161, in setup
    self._share_information_to_prevent_deadlock()
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 396, in _share_information_to_prevent_deadlock
    self._share_pids()
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 414, in _share_pids
    pids = self.all_gather(torch.tensor(os.getpid(), device=self.root_device))
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/strategies/parallel.py", line 113, in all_gather
    return all_gather_ddp_if_available(tensor, group=group, sync_grads=sync_grads)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/utilities/distributed.py", line 219, in all_gather_ddp_if_available
    return AllGatherGrad.apply(tensor, group)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/utilities/distributed.py", line 187, in forward
    torch.distributed.all_gather(gathered_tensor, tensor, group=group)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 2070, in all_gather
    work = group.allgather([tensor_list], [tensor])
RuntimeError: ProcessGroupGloo::allgather: unsupported device type mps

I'm looking to see if there's a way to turn off the distributed_backend=gloo

[tmm1]

Made some more progress by changing strategy from ddp to dp (https://pytorch-lightning.readthedocs.io/en/stable/extensions/strategy.html)

However, it seems ImageLogger which was using testtube is doing something important and switching to dp + csvlogger is not achieving the same result.

TypeError: LatentDiffusion.on_train_batch_start() missing 1 required positional argument: 'dataloader_idx'

EDIT: Found solution in Lightning-AI/pytorch-lightning#10315

[tmm1]

It is training!

Epoch 0: 11%|█▎ | 44/404 [03:59<32:37, 5.44s/it, loss=0.0784, v_num=0, train/loss_simple_step=0.00508, train/loss_vlb_step=2.81e-5, train/loss_step=0.00508, global_step=43.00]

I recall tho others saying training was broken on CUDA too in this fork, so I'm not sure if this is actually working or just appearing to. But a lot of the blockers are solved and we can get into the guts of the impl now.

EDIT: I am not seeing any of the warnings mentioned on the CUDA thread (related to batch_size)

[tmm1]

Died at the end.

Everything turned to nan at some point, I don't know if that's a bad sign.

I will try to make outputs optional and exit training early to see if it works.

Epoch 0: 100%|████████████████████| 404/404 [27:34<00:00,  4.10s/it, loss=nan, v_num=0, train/loss_simple_step=nan.0, train/loss_vlb_step=nan.0, train/loss_step=nan.0, global_step=399.0]


Summoning checkpoint.
Traceback (most recent call last):
  File "/Users/tmm1/code/stable-diffusion/./main.py", line 946, in <module>
    trainer.fit(model, data)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 696, in fit
    self._call_and_handle_interrupt(
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 650, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 735, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1166, in _run
    results = self._run_stage()
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1252, in _run_stage
    return self._run_train()
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1283, in _run_train
    self.fit_loop.run()
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/loops/loop.py", line 201, in run
    self.on_advance_end()
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 299, in on_advance_end
    self.trainer._call_callback_hooks("on_train_epoch_end")
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1597, in _call_callback_hooks
    fn(self, self.lightning_module, *args, **kwargs)
TypeError: CUDACallback.on_train_epoch_end() missing 1 required positional argument: 'outputs'

[tmm1]

Hmm got past last error but a new one now:

Epoch 0: 100%|███████████| 404/404 [16:24<00:00,  2.44s/it, loss=0.0835, v_num=0, train/loss_simple_step=0.00323, train/loss_vlb_step=1.88e-5, train/loss_step=0.00323, global_step=399.0/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:2075: LightningDeprecationWarning: `Trainer.root_gpu` is deprecated in v1.6 and will be removed in v1.8. Please use `Trainer.strategy.root_device.index` instead.
  rank_zero_deprecation(
Summoning checkpoint.
Traceback (most recent call last):
  File "/Users/tmm1/code/stable-diffusion/./main.py", line 946, in <module>
    trainer.fit(model, data)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 696, in fit
    self._call_and_handle_interrupt(
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 650, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 735, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1166, in _run
    results = self._run_stage()
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1252, in _run_stage
    return self._run_train()
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1283, in _run_train
    self.fit_loop.run()
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/loops/loop.py", line 201, in run
    self.on_advance_end()
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 299, in on_advance_end
    self.trainer._call_callback_hooks("on_train_epoch_end")
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1597, in _call_callback_hooks
    fn(self, self.lightning_module, *args, **kwargs)
  File "/Users/tmm1/code/stable-diffusion/main.py", line 558, in on_train_epoch_end
    torch.cuda.synchronize(trainer.root_gpu)
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/torch/cuda/__init__.py", line 494, in synchronize
    _lazy_init()
  File "/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/torch/cuda/__init__.py", line 211, in _lazy_init
    raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled

[tmm1]

Okay now its able to move onto epoch 1!

Average Epoch time: 472.13 seconds
Average Peak memory 0.00MiB
Epoch 1:  16%|█▌        | 63/404 [01:23<07:33,  1.33s/it, loss=0.0875, v_num=0, train/loss_simple_step=0.204, train/loss_vlb_step=0.000971, train/loss_step=0.204, global_step=462.0, train/loss_simple_epoch=0.108, train/loss_vlb_epoch=0.00118, train/loss_epoch=0.108]

[tmm1]

I have a checkpoints/embeddings_gs-1600.pt now but when I try using it, the output images are black :(

[tmm1]

I started fresh and by epoch 2 everything turns to nan. I think that is causing the black images?

Epoch 2: 26%|▎| 104/404 [02:11<06:20, 1.27s/it, loss=nan, v_num=0, train/loss_simple_step=nan.0, train/loss_vlb_step=nan.0, train/loss_step=nan.

cc @magnusviri @Birch-san

[Birch-san]

when I encountered black images with k-diffusion sampler, it was due to this problem (with ±Inf):
pytorch/pytorch#84364

fix was just to detach and clone the tensor:
crowsonkb/k-diffusion@3e976ef

if you're having NaN (rather than ±Inf), maybe that's unrelated.

I recommend to narrow down which line first introduces NaN. you can use this check to do so:

mycooltensor.isnan().any()
# returns a boolean

[tmm1]

Thanks @Birch-san! I see this warning at the start of training which may be related.

/opt/homebrew/anaconda3/envs/ldm/lib/python3.10/site-packages/pytorch_lightning/core/module.py:555: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
value = torch.tensor(value, device=self.device)