a-r-r-o-w / cogvideox-factory

Memory optimized finetuning scripts for CogVideoX using TorchAO and DeepSpeed
Apache License 2.0
237 stars 19 forks source link

Checkpoint recovery fails, checkpointing wont save #29

Open Cubey42 opened 3 days ago

Cubey42 commented 3 days ago
Resuming from checkpoint pytorch_model
Traceback (most recent call last):
  File "/home/cubeyai/cogvideox-factory/training/cogvideox_text_to_video_lora.py", line 924, in <module>
    main(args)
  File "/home/cubeyai/cogvideox-factory/training/cogvideox_text_to_video_lora.py", line 598, in main
    accelerator.load_state(os.path.join(args.output_dir, path))
  File "/home/cubeyai/cogvideox-factory/vvv/lib/python3.12/site-packages/accelerate/accelerator.py", line 3191, in load_state
    model.load_checkpoint(input_dir, ckpt_id, **load_model_func_kwargs)
  File "/home/cubeyai/cogvideox-factory/vvv/lib/python3.12/site-packages/deepspeed/runtime/engine.py", line 2815, in load_checkpoint
    load_path, client_states = self._load_checkpoint(load_dir,
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cubeyai/cogvideox-factory/vvv/lib/python3.12/site-packages/deepspeed/runtime/engine.py", line 2865, in _load_checkpoint
    sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine=self.checkpoint_engine)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cubeyai/cogvideox-factory/vvv/lib/python3.12/site-packages/deepspeed/runtime/state_dict_factory.py", line 43, in get_sd_loader
    return MegatronSDLoader(ckpt_list, version, checkpoint_engine)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cubeyai/cogvideox-factory/vvv/lib/python3.12/site-packages/deepspeed/runtime/state_dict_factory.py", line 193, in __init__
    super().__init__(ckpt_list, version, checkpoint_engine)
  File "/home/cubeyai/cogvideox-factory/vvv/lib/python3.12/site-packages/deepspeed/runtime/state_dict_factory.py", line 55, in __init__
    self.check_ckpt_list()
  File "/home/cubeyai/cogvideox-factory/vvv/lib/python3.12/site-packages/deepspeed/runtime/state_dict_factory.py", line 168, in check_ckpt_list
    assert len(self.ckpt_list) > 0
           ^^^^^^^^^^^^^^^^^^^^^^^

I had a training I was doing and when I set the checkpointing steps, it seems to fail, giving me only a model_states_00.pt , I tried to --resume_from_checkpoint and was met with this, I did see the note was #potentially but I was hoping to see if there was any hope to save this checkpoint before resorting to starting over.

The original issue with the checkpointing steps;

Steps: 100%|█| 3000/3000 [8:42:31<00:00, 10.17s/it, gradient_norm_after_clip=0, gradient_norm_before_clip=0,[2024-10-13 12:22:42,466] [INFO] [logging.py:96:log_dist] [Rank 0] [Torch] Checkpoint pytorch_model is about to be saved!
[2024-10-13 12:22:47,360] [INFO] [logging.py:96:log_dist] [Rank 0] Saving model checkpoint: ./results3/cogvideox-lora/checkpoint-3000/pytorch_model/mp_rank_00_model_states.pt
[2024-10-13 12:22:47,360] [INFO] [torch_checkpoint_engine.py:21:save] [Torch] Saving ./results3/cogvideox-lora/checkpoint-3000/pytorch_model/mp_rank_00_model_states.pt...
[2024-10-13 12:23:07,579] [INFO] [torch_checkpoint_engine.py:23:save] [Torch] Saved ./results3/cogvideox-lora/checkpoint-3000/pytorch_model/mp_rank_00_model_states.pt.
[2024-10-13 12:23:07,696] [INFO] [torch_checkpoint_engine.py:21:save] [Torch] Saving ./results3/cogvideox-lora/checkpoint-3000/pytorch_model/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt...
[2024-10-13 12:23:08,514] [INFO] [torch_checkpoint_engine.py:23:save] [Torch] Saved ./results3/cogvideox-lora/checkpoint-3000/pytorch_model/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt.
[2024-10-13 12:23:08,524] [INFO] [engine.py:3536:_save_zero_checkpoint] zero checkpoint saved ./results3/cogvideox-lora/checkpoint-3000/pytorch_model/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt
[2024-10-13 12:23:08,524] [INFO] [torch_checkpoint_engine.py:33:commit] [Torch] Checkpoint pytorch_model is ready now!
Traceback (most recent call last):
  File "/home/cubeyai/cogvideox-factory/training/cogvideox_text_to_video_lora.py", line 924, in <module>
    main(args)
  File "/home/cubeyai/cogvideox-factory/training/cogvideox_text_to_video_lora.py", line 755, in main
    accelerator.save_state(save_path)
  File "/home/cubeyai/cogvideox-factory/vvv/lib/python3.12/site-packages/accelerate/accelerator.py", line 3081, in save_state
    hook(self._models, weights, output_dir)
  File "/home/cubeyai/cogvideox-factory/training/cogvideox_text_to_video_lora.py", line 363, in save_model_hook
    raise ValueError(f"unexpected save model: {model.__class__}")
ValueError: unexpected save model: <class 'deepspeed.runtime.engine.DeepSpeedEngine'>
sayakpaul commented 22 hours ago

@Cubey42 could you give https://github.com/a-r-r-o-w/cogvideox-factory/pull/38 a try?