fsdp when training?

[tongdaxu]

Hi,

Thank you for your great project.

I am running the pokemon demo with fsdp=True, 8xA30 GPU, I have the following configs

use_fsdp=False   # if use FSDP mode
num_workers=10
train_batch_size = 64 # 32
num_epochs = 200 # 3
gradient_accumulation_steps = 2
grad_checkpointing = True
gradient_clip = 0.01
fp32_attention = True

And it turns out that

Traceback (most recent call last):
  File "/NEW_EDS/JJ_Group/xutd/PixArt-alpha/train.py", line 333, in <module>
    train()
  File "/NEW_EDS/JJ_Group/xutd/PixArt-alpha/train.py", line 90, in train
    loss_term = train_diffusion.training_losses(model, clean_images, timesteps, model_kwargs=dict(y=y, mask=y_mask, data_info=data_info))
  File "/NEW_EDS/JJ_Group/xutd/PixArt-alpha/diffusion/model/respace.py", line 94, in training_losses
    return super().training_losses(self._wrap_model(model), *args, **kwargs)
  File "/NEW_EDS/JJ_Group/xutd/PixArt-alpha/diffusion/model/gaussian_diffusion.py", line 750, in training_losses
    model_output = model(x_t, t, **model_kwargs)
  File "/NEW_EDS/JJ_Group/xutd/PixArt-alpha/diffusion/model/respace.py", line 131, in __call__
    return self.model(x, timestep=new_ts, **kwargs)
  File "/home/xutd/.conda/envs/dvp2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/xutd/.conda/envs/dvp2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/xutd/.conda/envs/dvp2/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 839, in forward
    output = self._fsdp_wrapped_module(*args, **kwargs)
  File "/home/xutd/.conda/envs/dvp2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/xutd/.conda/envs/dvp2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/xutd/.conda/envs/dvp2/lib/python3.10/site-packages/accelerate/utils/operations.py", line 687, in forward
    return model_forward(*args, **kwargs)
  File "/home/xutd/.conda/envs/dvp2/lib/python3.10/site-packages/accelerate/utils/operations.py", line 675, in __call__
    return convert_to_fp32(self.model_forward(*args, **kwargs))
  File "/home/xutd/.conda/envs/dvp2/lib/python3.10/site-packages/torch/amp/autocast_mode.py", line 16, in decorate_autocast
    return func(*args, **kwargs)
  File "/NEW_EDS/JJ_Group/xutd/PixArt-alpha/diffusion/model/nets/PixArt.py", line 121, in forward
    t = self.t_embedder(timestep.to(x.dtype))  # (N, D)
  File "/home/xutd/.conda/envs/dvp2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/xutd/.conda/envs/dvp2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/NEW_EDS/JJ_Group/xutd/PixArt-alpha/diffusion/model/nets/PixArt_blocks.py", line 269, in forward
    t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(self.dtype)
  File "/NEW_EDS/JJ_Group/xutd/PixArt-alpha/diffusion/model/nets/PixArt_blocks.py", line 275, in dtype
    return next(self.parameters()).dtype

I have tried to turn fp32_attention = False but it remains the same.

Any ideas why?

I am trying to seqeeze more batch size from 24 GB memory, and I would love your advices.

Thank you.

[lawrence-cj]

We haven't tested the FSDP version. Will look into it later.