[BUG]: can not find parameter `{full_param_name}` in the GeminiDDP module

[FrankieDong]

🐛 Describe the bug

When I run the train code, it occurs:

Traceback (most recent call last):
File "/home/zrytest/colorSD2.0/ColossalAI-main/ColossalAI/examples/images/diffusion/main.py", line 804, in
trainer.fit(model, data)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 602, in fit
call._call_and_handle_interrupt(
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/call.py", line 38, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 644, in _fit_impl
self._run(model, ckpt_path=self.ckpt_path)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1097, in _run
results = self._run_stage()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1176, in _run_stage
self._run_train()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1199, in _run_train
Traceback (most recent call last):
File "/home/zrytest/colorSD2.0/ColossalAI-main/ColossalAI/examples/images/diffusion/main.py", line 804, in
self.fit_loop.run()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/loops/loop.py", line 200, in run
self.on_advance_end()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/loops/fit_loop.py", line 295, in on_advance_end
self.trainer._call_callback_hooks("on_train_epoch_end")
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1379, in _call_callback_hooks
trainer.fit(model, data)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 602, in fit
call._call_and_handle_interrupt(
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/call.py", line 38, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 644, in _fit_impl
fn(self, self.lightning_module, *args, **kwargs)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 304, in on_train_epoch_end
self._save_last_checkpoint(trainer, monitor_candidates)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 649, in _save_last_checkpoint
self._run(model, ckpt_path=self.ckpt_path)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1097, in _run
self._save_checkpoint(trainer, filepath)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 365, in _save_checkpoint
results = self._run_stage()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1176, in _run_stage
trainer.save_checkpoint(filepath, self.save_weights_only)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1924, in save_checkpoint
self._run_train()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1199, in _run_train
self._checkpoint_connector.save_checkpoint(filepath, weights_only=weights_only, storage_options=storage_options)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/checkpoint_connector.py", line 510, in save_checkpoint
self.fit_loop.run()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/loops/loop.py", line 200, in run
_checkpoint = self.dump_checkpoint(weights_only)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/checkpoint_connector.py", line 442, in dump_checkpoint
self.on_advance_end()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/loops/fit_loop.py", line 295, in on_advance_end
"state_dict": self._get_lightning_module_state_dict(),
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/checkpoint_connector.py", line 524, in _get_lightning_module_state_dict
self.trainer._call_callback_hooks("on_train_epoch_end")
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1379, in _call_callback_hooks
state_dict = self.trainer.strategy.lightning_module_state_dict()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/strategies/colossalai.py", line 428, in lightning_module_state_dict
org_dict = self.model.state_dict(only_rank_0=rank_zero_only)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/colossalai-0.2.0-py3.8.egg/colossalai/nn/parallel/data_parallel.py", line 352, in state_dict
fn(self, self.lightning_module, *args, **kwargs)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 304, in on_train_epoch_end
torch_model = get_static_torch_model(zero_ddp_model=self, only_rank_0=only_rank_0)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/colossalai-0.2.0-py3.8.egg/colossalai/nn/parallel/utils.py", line 101, in get_static_torch_model
self._save_last_checkpoint(trainer, monitor_candidates)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 649, in _save_last_checkpoint
assert param in colo_to_torch, f"can not find parameter {full_param_name} in the GeminiDDP module"
AssertionError: can not find parameter _forward_module.first_stage_model.encoder.conv_in.weight in the GeminiDDP module

Environment

No response

[FrankieDong]

lightning used stable/colossalai
the same problem occured in docker env and physic env

[1SAA]

Hi @FrankieDong

This seems like a bug caused by ignored ddp parameters. Could you tell me how to repeat your this bug? So that I can check whether this bug is fixed.

[FrankieDong]

Of course
First I git clone the main branch in the github.
Second I use the "train_colossalai_cifar10.yaml" without editing
Third I run the default command in the "example/image/diffusion"

then occured this error

[FrankieDong]

{'accelerator': 'gpu', 'devices': 1, 'log_gpu_memory': 'all', 'max_epochs': 2, 'precision': 16, 'auto_select_gpus': False, 'strategy': {'target': 'strategies.ColossalAIStrategy', 'params': {'use_chunk': True, 'enable_distributed_storage': True, 'placement_policy': 'cuda', 'force_outputs_fp32': True}}, 'log_every_n_steps': 2, 'logger': True, 'default_root_dir': '/tmp/diff_log/'}

[Issues-translate-bot]

Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿

---

{'accelerator': 'gpu', 'devices': 1, 'log_gpu_memory': 'all', 'max_epochs': 2, 'precision': 16, 'auto_select_gpus': False, 'strategy': {'target': 'strategies.ColossalAIStrategy', 'params': {'use_chunk': True, 'enable_distributed_storage': True, 'placement_policy': 'cuda', 'force_outputs_fp32': True}}, 'log_every_n_steps': 2, 'logger': True, 'default_root_dir': '/tmp/diff_log/'}

[FrankieDong]

Epoch 0: 100%|█████████████████████████████████████████████████████████| 1/1 [00:13<00:00, 13.14s/it, loss=0.436, v_num=0, train/loss_simple_step=0.436, train/loss_vlb_step=0.436, train/loss_step=0.436, global_step=0.000, lr_abs=2.56e-8]/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/torch/optim/lr_scheduler.py:131: UserWarning: Detected call of lr_scheduler.step() before optimizer.step(). In PyTorch 1.1.0 and later, you should call them in the opposite order: optimizer.step() before lr_scheduler.step(). Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
warnings.warn("Detected call of lr_scheduler.step() before optimizer.step(). "
/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/torch/optim/lr_scheduler.py:131: UserWarning: Detected call of lr_scheduler.step() before optimizer.step(). In PyTorch 1.1.0 and later, you should call them in the opposite order: optimizer.step() before lr_scheduler.step(). Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
warnings.warn("Detected call of lr_scheduler.step() before optimizer.step(). "
/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/torch/optim/lr_scheduler.py:131: UserWarning: Detected call of lr_scheduler.step() before optimizer.step(). In PyTorch 1.1.0 and later, you should call them in the opposite order: optimizer.step() before lr_scheduler.step(). Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
warnings.warn("Detected call of lr_scheduler.step() before optimizer.step(). "
/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py:536: PossibleUserWarning: It is recommended to use self.log('train/loss_simple', ..., sync_dist=True) when logging on epoch level in distributed setting to accumulate the metric across devices.
warning_cache.warn(
/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py:536: PossibleUserWarning: It is recommended to use self.log('train/loss_vlb', ..., sync_dist=True) when logging on epoch level in distributed setting to accumulate the metric across devices.
warning_cache.warn(
/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py:536: PossibleUserWarning: It is recommended to use self.log('train/loss', ..., sync_dist=True) when logging on epoch level in distributed setting to accumulate the metric across devices.
warning_cache.warn(
Epoch 0: 100%|█| 1/1 [00:13<00:00, 13.15s/it, loss=0.436, v_num=0, train/loss_simple_step=0.436, train/loss_vlb_step=0.436, train/loss_step=0.436, global_step=0.000, lr_abs=2.56e-8, train/loss_simple_epoch=0.436, train/loss_vlb_epoch=0.4/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py:359: UserWarning: ModelCheckpoint(monitor='val/loss_simple_ema') could not find the monitored key in the returned metrics: ['train/loss_simple', 'train/loss_simple_step', 'train/loss_vlb', 'train/loss_vlb_step', 'train/loss', 'train/loss_step', 'global_step', 'lr_abs', 'train/loss_simple_epoch', 'train/loss_vlb_epoch', 'train/loss_epoch', 'epoch', 'step']. HINT: Did you call log('val/loss_simple_ema', value) in the LightningModule?
warning_cache.warn(m)
/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/strategies/ddp.py:437: UserWarning: Error handling mechanism for deadlock detection is uninitialized. Skipping check.
rank_zero_warn("Error handling mechanism for deadlock detection is uninitialized. Skipping check.")
Summoning checkpoint.
Traceback (most recent call last):
File "/home/zrytest/colorSD2.0/ColossalAI-main/ColossalAI/examples/images/diffusion/main.py", line 804, in
Traceback (most recent call last):
File "/home/zrytest/colorSD2.0/ColossalAI-main/ColossalAI/examples/images/diffusion/main.py", line 804, in
trainer.fit(model, data)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 602, in fit
trainer.fit(model, data)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 602, in fit
call._call_and_handle_interrupt(
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/call.py", line 38, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 644, in _fit_impl
call._call_and_handle_interrupt(
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/call.py", line 38, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 644, in _fit_impl
self._run(model, ckpt_path=self.ckpt_path)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1097, in _run
self._run(model, ckpt_path=self.ckpt_path)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1097, in _run
results = self._run_stage()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1176, in _run_stage
results = self._run_stage()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1176, in _run_stage
self._run_train()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1199, in _run_train
self._run_train()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1199, in _run_train
self.fit_loop.run()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/loops/loop.py", line 200, in run
self.on_advance_end()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/loops/fit_loop.py", line 295, in on_advance_end
self.trainer._call_callback_hooks("on_train_epoch_end")
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1379, in _call_callback_hooks
self.fit_loop.run()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/loops/loop.py", line 200, in run
self.on_advance_end()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/loops/fit_loop.py", line 295, in on_advance_end
self.trainer._call_callback_hooks("on_train_epoch_end")
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1379, in _call_callback_hooks
fn(self, self.lightning_module, *args, **kwargs)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 304, in on_train_epoch_end
self._save_last_checkpoint(trainer, monitor_candidates)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 649, in _save_last_checkpoint
self._save_checkpoint(trainer, filepath)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 365, in _save_checkpoint
fn(self, self.lightning_module, *args, **kwargs)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 304, in on_train_epoch_end
trainer.save_checkpoint(filepath, self.save_weights_only)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1924, in save_checkpoint
self._save_last_checkpoint(trainer, monitor_candidates)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 649, in _save_last_checkpoint
self._save_checkpoint(trainer, filepath)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 365, in _save_checkpoint
trainer.save_checkpoint(filepath, self.save_weights_only)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1924, in save_checkpoint
self._checkpoint_connector.save_checkpoint(filepath, weights_only=weights_only, storage_options=storage_options)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/checkpoint_connector.py", line 510, in save_checkpoint
_checkpoint = self.dump_checkpoint(weights_only)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/checkpoint_connector.py", line 442, in dump_checkpoint
"state_dict": self._get_lightning_module_state_dict(),
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/checkpoint_connector.py", line 524, in _get_lightning_module_state_dict
self._checkpoint_connector.save_checkpoint(filepath, weights_only=weights_only, storage_options=storage_options)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/checkpoint_connector.py", line 510, in save_checkpoint
state_dict = self.trainer.strategy.lightning_module_state_dict()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/strategies/colossalai.py", line 428, in lightning_module_state_dict
org_dict = self.model.state_dict(only_rank_0=rank_zero_only)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/colossalai-0.2.0-py3.8.egg/colossalai/nn/parallel/data_parallel.py", line 352, in state_dict
_checkpoint = self.dump_checkpoint(weights_only)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/checkpoint_connector.py", line 442, in dump_checkpoint
torch_model = get_static_torch_model(zero_ddp_model=self, only_rank_0=only_rank_0)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/colossalai-0.2.0-py3.8.egg/colossalai/nn/parallel/utils.py", line 101, in get_static_torch_model
"state_dict": self._get_lightning_module_state_dict(),
assert param in colo_to_torch, f"can not find parameter {full_param_name} in the GeminiDDP module"
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/checkpoint_connector.py", line 524, in _get_lightning_module_state_dict
AssertionError: can not find parameter _forward_module.first_stage_model.encoder.conv_in.weight in the GeminiDDP module
state_dict = self.trainer.strategy.lightning_module_state_dict()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/strategies/colossalai.py", line 428, in lightning_module_state_dict
org_dict = self.model.state_dict(only_rank_0=rank_zero_only)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/colossalai-0.2.0-py3.8.egg/colossalai/nn/parallel/data_parallel.py", line 352, in state_dict
torch_model = get_static_torch_model(zero_ddp_model=self, only_rank_0=only_rank_0)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/colossalai-0.2.0-py3.8.egg/colossalai/nn/parallel/utils.py", line 101, in get_static_torch_model
assert param in colo_to_torch, f"can not find parameter {full_param_name} in the GeminiDDP module"
AssertionError: can not find parameter _forward_module.first_stage_model.encoder.conv_in.weight in the GeminiDDP module
Traceback (most recent call last):
File "/home/zrytest/colorSD2.0/ColossalAI-main/ColossalAI/examples/images/diffusion/main.py", line 804, in
trainer.fit(model, data)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 602, in fit
call._call_and_handle_interrupt(
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/call.py", line 38, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 644, in _fit_impl
self._run(model, ckpt_path=self.ckpt_path)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1097, in _run
results = self._run_stage()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1176, in _run_stage
self._run_train()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1199, in _run_train
self.fit_loop.run()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/loops/loop.py", line 200, in run
self.on_advance_end()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/loops/fit_loop.py", line 295, in on_advance_end
self.trainer._call_callback_hooks("on_train_epoch_end")
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1379, in _call_callback_hooks
fn(self, self.lightning_module, *args, **kwargs)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 304, in on_train_epoch_end
self._save_last_checkpoint(trainer, monitor_candidates)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 649, in _save_last_checkpoint
self._save_checkpoint(trainer, filepath)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/callbacks/model_checkpoint.py", line 365, in _save_checkpoint
trainer.save_checkpoint(filepath, self.save_weights_only)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/trainer.py", line 1924, in save_checkpoint
self._checkpoint_connector.save_checkpoint(filepath, weights_only=weights_only, storage_options=storage_options)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/checkpoint_connector.py", line 510, in save_checkpoint
_checkpoint = self.dump_checkpoint(weights_only)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/checkpoint_connector.py", line 442, in dump_checkpoint
"state_dict": self._get_lightning_module_state_dict(),
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/trainer/connectors/checkpoint_connector.py", line 524, in _get_lightning_module_state_dict
state_dict = self.trainer.strategy.lightning_module_state_dict()
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/lightning/pytorch/strategies/colossalai.py", line 428, in lightning_module_state_dict
org_dict = self.model.state_dict(only_rank_0=rank_zero_only)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/colossalai-0.2.0-py3.8.egg/colossalai/nn/parallel/data_parallel.py", line 352, in state_dict
torch_model = get_static_torch_model(zero_ddp_model=self, only_rank_0=only_rank_0)
File "/home/zrytest/anaconda3/envs/ldm_dx/lib/python3.8/site-packages/colossalai-0.2.0-py3.8.egg/colossalai/nn/parallel/utils.py", line 101, in get_static_torch_model
assert param in colo_to_torch, f"can not find parameter {full_param_name} in the GeminiDDP module"
AssertionError: can not find parameter _forward_module.first_stage_model.encoder.conv_in.weight in the GeminiDDP module

[FrankieDong]

Hi @FrankieDong

This seems like a bug caused by ignored ddp parameters. Could you tell me how to repeat your this bug? So that I can check whether this bug is fixed.

I upload more info, please check, 3ks