resume from checkpoint fails [Krea 2]

[paju1986]

Hi.
When i try to resume from checkpoint my Krea2 trainig i get this error:

WARNING:root:Missing weight for layer last.linear
Resuming training from checkpoint. Resuming at epoch: 21, step: 3641
[rank0]: Traceback (most recent call last):
[rank0]: File "/mnt/AI/diffusion-pipe/train.py", line 908, in
[rank0]: loss = model_engine.train_batch(iterator).item()
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/deepspeed/runtime/pipe/engine.py", line 388, in train_batch
[rank0]: self._exec_schedule(sched)
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/deepspeed/runtime/pipe/engine.py", line 1368, in _exec_schedule
[rank0]: self._exec_instr(**cmd.kwargs)
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/deepspeed/runtime/pipe/engine.py", line 737, in _exec_forward_pass
[rank0]: outputs = super().forward(inputs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank0]: ret_val = func(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/deepspeed/runtime/engine.py", line 2236, in forward
[rank0]: loss = self.module(*inputs, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl
[rank0]: return inner()
[rank0]: ^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner
[rank0]: result = forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/deepspeed/runtime/pipe/module.py", line 388, in forward
[rank0]: x = self.activation_checkpoint_func(exec_range_func(start_idx, end_idx), *x)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/torch/_compile.py", line 54, in inner
[rank0]: return disable_fn(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 1298, in _fn
[rank0]: return fn(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/torch/utils/checkpoint.py", line 518, in checkpoint
[rank0]: ret = function(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/deepspeed/runtime/pipe/module.py", line 365, in exec_func
[rank0]: inputs = layer(inputs)
[rank0]: ^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
[rank0]: return func(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/AI/diffusion-pipe/models/krea2.py", line 211, in forward
[rank0]: img = self.first(img)
[rank0]: ^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
[rank0]: return self._call_impl(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/home/alberto/.conda/envs/diffusion-pipeRocm/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
[rank0]: return forward_call(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/AI/diffusion-pipe/utils/../submodules/ComfyUI/comfy/ops.py", line 1244, in forward
[rank0]: output = self.forward_comfy_cast_weights(input, compute_dtype, want_requant=isinstance(input, QuantizedTensor))
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/AI/diffusion-pipe/utils/../submodules/ComfyUI/comfy/ops.py", line 1187, in forward_comfy_cast_weights
[rank0]: weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True, compute_dtype=compute_dtype, want_requant=want_requant)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/mnt/AI/diffusion-pipe/utils/../submodules/ComfyUI/comfy/ops.py", line 337, in cast_bias_weight
[rank0]: if offloadable and (device != s.weight.device or
[rank0]: ^^^^^^^^^^^^^^^
[rank0]: AttributeError: 'NoneType' object has no attribute 'device'
[2026-06-29 09:52:35,376] [INFO] [launch.py:335:sigkill_handler] Killing subprocess 42979
[2026-06-29 09:52:35,376] [ERROR] [launch.py:341:sigkill_handler] ['/home/alberto/.conda/envs/diffusion-pipeRocm/bin/python3.12', '-u', 'train.py', '--local_rank=0', '--deepspeed', '--resume_from_checkpoint', '20260628_18-57-27', '--config', './krea2.toml'] exits with return code = 1

The training went well, but i wanted to keep training a little more but i get this. Im using rocm 7.2

Regards

[paju1986]

any news?

[tdrussell]

Resuming from checkpoint works for me with Krea2. ROCm is likely the issue. The error is coming from inside the ComfyUI code but I don't know why. I probably can't help much without being able to reproduce the error myself and debug it.

[paju1986]

strange it used to work fine with other models maybe was after updating the submodules