Got errors when initiating training

[ThisisBillhe]

[rank6]:     self.ema_model = EMAModel(self.accelerator.unwrap_model(self.model), self.device)
[rank6]:                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank6]:   File "/mnt/workspace/miniconda3/envs/blip3o/lib/python3.11/site-packages/accelerate/accelerator.py", line 2866, in unwrap_model
[rank6]:     return extract_model_from_parallel(model, keep_fp32_wrapper, keep_torch_compile)
[rank6]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank6]:   File "/mnt/workspace/miniconda3/envs/blip3o/lib/python3.11/site-packages/accelerate/utils/other.py", line 180, in extract_model_from_parallel
[rank6]:     model = model._orig_mod
[rank6]:             ^^^^^^^^^^^^^^^
[rank6]:   File "/mnt/workspace/miniconda3/envs/blip3o/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1931, in __getattr__
[rank6]:     raise AttributeError(
[rank6]: AttributeError: 'DistributedDataParallel' object has no attribute '_orig_mod'

While this error seems to be related to DDP and accelerate, I can barely find information on the Google. Do you have any idea how to address this?

[ThisisBillhe]

I think the reason is that model.vae/dit/encodr have been compiled, while the model itself has not. So that accelerate library thinks there is _orig_mod attribute.

[DTennant]

Hi, can you try turn off enable_ema and see if that resolves the issue? https://github.com/visual-gen/semanticist/blob/main/configs/tokenizer_l.yaml#L15

[ThisisBillhe]

Hi Bingchen, thanks for your reply. Will this affect the final performance?

[xwen99]

Yes, turning off EMA could affect performance. If you prefer to keep enabling torch.compile and EMA, I would suggest either 1) add a new function to unwrap the model manually before initializing EMA, or 2) simply feed the EMA model with self.model, and add a script to unwrap the final EMA weights (cleanup ._orig_mod and .module prefixes) after training is done. We used the 2nd solution while developing this project and are happy to see better workarounds/PRs.