finetune OOM evalue lm_loss, metrics = forward_step(data_iterator, model, args, timers)

[sunshinejuanjuan123]

run cmd in my data to finetune

bash finetune_single_identity.sh # Multi GPUs

in 8GPUS H100 eval will be OOM

from sat.training.deepspeed_training import training_main

    training_main(

        args,

        model_cls=SATVideoDiffusionEngine,

        forward_step_function=partial(forward_step, data_class=data_class),

        forward_step_eval=partial(

            forward_step_eval, data_class=data_class, only_log_video_latents=args.only_log_video_latents

        ),

        create_dataset_function=create_dataset_function,

    )

sat.training.deepspeed_training import training_main eval时显存会OOM如何修复解决
从 sat.training.deepspeed_training 导入 training_main

training_main(

args,

model_cls=SATVideoDiffusionEngine,

forward_step_function=partial(forward_step, data_class=data_class),

forward_step_eval=partial(

forward_step_eval, data_class=data_class, only_log_video_latents=args.only_log_video_latents

        ),

create_dataset_function=create_dataset_function,

    )

sat.training.deepspeed_training 导入 training_main 评估时显存会 OOM 如何修复解决

/mnt/afs//Concat-ID/train_video.py:74: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.                                    
  with torch.no_grad(), torch.cuda.amp.autocast(**gpu_autocast_kwargs):                                                                                                                                          
[rank0]: Traceback (most recent call last):                                                                                                                                                                      
[rank0]:   File "/mnt/afs//Concat-ID/train_video.py", line 235, in <module>                                                                                                                          
[rank0]:     training_main(                                                                                                                                                                                      
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/training/deepspeed_training.py", line 158, in training_main                                                                   
[rank0]:     iteration, skipped = train(model, optimizer,                                                                                                                                                        
[rank0]:                          ^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                        
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/training/deepspeed_training.py", line 414, in train                                                                           
[rank0]:     evaluate_and_print_results(                                                                                                                                                                         
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/training/deepspeed_training.py", line 629, in evaluate_and_print_results                                                      
[rank0]:     lm_loss, metrics = evaluate(data_iterator, model, eval_iters, args, timers, split, verbose, has_last, hooks=hooks)                                                                                  
[rank0]:                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                  
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/training/deepspeed_training.py", line 571, in evaluate                                                                        
[rank0]:     lm_loss, metrics = forward_step(data_iterator, model, args, timers)                                                                                                                                 
[rank0]:                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                 
[rank0]:   File "/mnt/afs//Concat-ID/train_video.py", line 173, in forward_step_eval                                                                                                                 
[rank0]:     log_video(batch_video, model, args, only_log_video_latents=only_log_video_latents)                                                                                                                  
[rank0]:   File "/mnt/afs//Concat-ID/train_video.py", line 75, in log_video                                                                                                                          
[rank0]:     videos = model.log_video(batch, only_log_video_latents=only_log_video_latents)                                                                                                                      
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                      
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context                                                                        
[rank0]:     return func(*args, **kwargs)                                                                                                                                                                        
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                                        
[rank0]:   File "/mnt/afs//Concat-ID/diffusion_video.py", line 334, in log_video                                                                                                                     
[rank0]:     log["reconstructions"] = self.decode_first_stage(z).to(torch.float32)                   

[rank0]:            ^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                               [45/1891]
[rank0]:   File "/mnt/afs//Concat-ID/diffusion_video.py", line 334, in log_video                                                                                                                     
[rank0]:     log["reconstructions"] = self.decode_first_stage(z).to(torch.float32)                                                                                                                               
[rank0]:                              ^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                 
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context                                                                        
[rank0]:     return func(*args, **kwargs)                                                                                                                                                                        
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                                        
[rank0]:   File "/mnt/afs//Concat-ID/diffusion_video.py", line 203, in decode_first_stage                                                                                                            
[rank0]:     out = self.first_stage_model.decode(z[n * n_samples : (n + 1) * n_samples], **kwargs)                                                                                                               
[rank0]:           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                               
[rank0]:   File "/mnt/afs//Concat-ID/vae_modules/autoencoder.py", line 634, in decode                                                                                                                
[rank0]:     x = super().decode(z, **kwargs)                                                                                                                                                                     
[rank0]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                                     
[rank0]:   File "/mnt/afs//Concat-ID/vae_modules/autoencoder.py", line 233, in decode                                                                                                                
[rank0]:     x = self.decoder(z, **kwargs)                                                                                                                                                                       
[rank0]:         ^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                                       
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl                                                                     
[rank0]:     return self._call_impl(*args, **kwargs)                                                                                                                                                             
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                             
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl                                                                             
[rank0]:     return forward_call(*args, **kwargs)                                                                                                                                                                
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                                
[rank0]:   File "/mnt/afs//Concat-ID/vae_modules/cp_enc_dec.py", line 970, in forward                                                                                                                
[rank0]:     h = self.up[i_level].block[i_block](h, temb, zq, clear_fake_cp_cache=clear_fake_cp_cache)                                                                                                           
[rank0]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                           
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl                                                                     
[rank0]:     return self._call_impl(*args, **kwargs)                                                                                                                                                             
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                             
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl                                                                             
[rank0]:     return forward_call(*args, **kwargs)                                                                                                                                                                
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                                
[rank0]:   File "/mnt/afs//Concat-ID/vae_modules/cp_enc_dec.py", line 682, in forward                                                                                                                
[rank0]:     h = self.norm1(h, zq, clear_fake_cp_cache=clear_fake_cp_cache)                                                                                                                                      
[rank0]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                      
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl                                                                     
[rank0]:     return self._call_impl(*args, **kwargs)                                                                                                                                                             
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                             
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl                                                                             
[rank0]:     return forward_call(*args, **kwargs)                                                                                                                                                                
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                                
[rank0]:   File "/mnt/afs//Concat-ID/vae_modules/cp_enc_dec.py", line 509, in forward                                                                                                                
[rank0]:     new_f = norm_f * self.conv_y(zq) + self.conv_b(zq)                                                                                                                                                  
[rank0]:             ~~~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~                                                                                                                                                  
[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 8.08 GiB. GPU 0 has a total capacity of 79.33 GiB of which 1.54 GiB is free. Process 293276 has 77.75 GiB memory in use. Of the allocated 
memory 74.27 GiB is allocated by PyTorch, and 108.13 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoi
d fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)  

[rank0]: File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/training/deepspeed_training.py", line 571, in evaluate
[rank0]: lm_loss, metrics = forward_step(data_iterator, model, args, timers)
尝试在此处加torch.cuda.empty_cache() gc.collect() 仍然会出现OOM

[sunshinejuanjuan123]

修改zero_optimization
stage: 3
报错

[2025-04-12 09:52:21,831] [INFO] [partition_parameters.py:348:__exit__] finished initializing model - num_params = 4, num_elems = 0.00B                                                                [218/1988]
[rank2]: Traceback (most recent call last):                                                                                                                                                                      
[rank2]:   File "/mnt/afs//Concat-ID/train_video.py", line 235, in <module>                                                                                                                          
[rank2]:     training_main(                                                                                                                                                                                      
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/training/deepspeed_training.py", line 84, in training_main                                                                    
[rank2]:     model = get_model(args, model_cls)                                                                                                                                                                  
[rank2]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                                  
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/base_model.py", line 418, in get_model                                                                                  
[rank2]:     model = model_cls(args, params_dtype=params_dtype, **kwargs)                                                                                                                                        
[rank2]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                        
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper                                                             
[rank2]:     f(module, *args, **kwargs)                                                                                                                                                                          
[rank2]:   File "/mnt/afs//Concat-ID/diffusion_video.py", line 74, in __init__                                                                                                                       
[rank2]:     model = instantiate_from_config(network_config)                                                                                                                                                     
[rank2]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                     
[rank2]:   File "/mnt/afs//Concat-ID/sgm/util.py", line 261, in instantiate_from_config                                                                                                              
[rank2]:     return get_obj_from_str(config["target"])(**config.get("params", dict()), **extra_kwargs)                                                                                                           
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                           
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper                                                             
[rank2]:     f(module, *args, **kwargs)                                                                                                                                                                          
[rank2]:   File "/mnt/afs//Concat-ID/dit_video_concat.py", line 662, in __init__                                                                                                                     
[rank2]:     super().__init__(args=transformer_args, transformer=None, **kwargs)
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank2]:     f(module, *args, **kwargs)
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/base_model.py", line 93, in __init__
[rank2]:     self.transformer = BaseTransformer(
[rank2]:                        ^^^^^^^^^^^^^^^^
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank2]:     f(module, *args, **kwargs)
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 565, in __init__
[rank2]:     [get_layer(layer_id) for layer_id in range(num_layers)])
[rank2]:     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 565, in <listcomp>
[rank2]:     [get_layer(layer_id) for layer_id in range(num_layers)])
[rank2]:      ^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 532, in get_layer
[rank2]:     return BaseTransformerLayer(
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank2]:     f(module, *args, **kwargs)
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 363, in __init__
[rank2]:     self.attention = SelfAttention(
[rank2]:                      ^^^^^^^^^^^^^^
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank2]:     f(module, *args, **kwargs)
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 66, in __init__
[rank2]:     self.query_key_value = ColumnParallelLinear(
[rank2]:                            ^^^^^^^^^^^^^^^^^^^^^
                                                                                                                                    [172/1988]
[rank2]:                            ^^^^^^^^^^^^^^^^^^^^^
[rank2]: TypeError: ColumnParallelLinear.__init__() got multiple values for argument 'module'
[rank0]: Traceback (most recent call last):
[rank0]:   File "/mnt/afs//Concat-ID/train_video.py", line 235, in <module>
[rank0]:     training_main(
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/training/deepspeed_training.py", line 84, in training_main
[rank0]:     model = get_model(args, model_cls)
[rank0]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/base_model.py", line 418, in get_model
[rank0]:     model = model_cls(args, params_dtype=params_dtype, **kwargs)
[rank0]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank0]:     f(module, *args, **kwargs)
[rank0]:   File "/mnt/afs//Concat-ID/diffusion_video.py", line 74, in __init__
[rank0]:     model = instantiate_from_config(network_config)
[rank0]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/afs//Concat-ID/sgm/util.py", line 261, in instantiate_from_config
[rank0]:     return get_obj_from_str(config["target"])(**config.get("params", dict()), **extra_kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank0]:     f(module, *args, **kwargs)
[rank0]:   File "/mnt/afs//Concat-ID/dit_video_concat.py", line 662, in __init__
[rank0]:     super().__init__(args=transformer_args, transformer=None, **kwargs)
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank0]:     f(module, *args, **kwargs)
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/base_model.py", line 93, in __init__
[rank0]:     self.transformer = BaseTransformer(
[rank0]:                        ^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank0]:     f(module, *args, **kwargs)
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 565, in __init__
[rank0]:     [get_layer(layer_id) for layer_id in range(num_layers)])
[rank0]:     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 565, in <listcomp>
[rank0]:     [get_layer(layer_id) for layer_id in range(num_layers)])
[rank0]:      ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 532, in get_layer
[rank0]:     return BaseTransformerLayer(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank0]:     f(module, *args, **kwargs)
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 363, in __init__
[rank0]:     self.attention = SelfAttention(
[rank0]:                      ^^^^^^^^^^^^^^
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank0]:     f(module, *args, **kwargs)
[rank0]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 66, in __init__

                                                                             
[rank0]:     self.query_key_value = ColumnParallelLinear(                                                                                                                                                        
[rank0]:                            ^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                        
[rank0]: TypeError: ColumnParallelLinear.__init__() got multiple values for argument 'module'                                                                                                                    
[rank1]: Traceback (most recent call last):
[rank1]:   File "/mnt/afs//Concat-ID/train_video.py", line 235, in <module>
[rank1]:     training_main(
[rank1]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/training/deepspeed_training.py", line 84, in training_main
[rank1]:     model = get_model(args, model_cls)
[rank1]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/base_model.py", line 418, in get_model
[rank1]:     model = model_cls(args, params_dtype=params_dtype, **kwargs)
[rank1]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank1]:     f(module, *args, **kwargs)
[rank1]:   File "/mnt/afs//Concat-ID/diffusion_video.py", line 74, in __init__
[rank1]:     model = instantiate_from_config(network_config)
[rank1]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/afs//Concat-ID/sgm/util.py", line 261, in instantiate_from_config
[rank1]:     return get_obj_from_str(config["target"])(**config.get("params", dict()), **extra_kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank1]:     f(module, *args, **kwargs)
[rank1]:   File "/mnt/afs//Concat-ID/dit_video_concat.py", line 662, in __init__
[rank1]:     super().__init__(args=transformer_args, transformer=None, **kwargs)
[rank1]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank1]:     f(module, *args, **kwargs)
[rank1]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/base_model.py", line 93, in __init__
[rank1]:     self.transformer = BaseTransformer(
[rank1]:                        ^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank1]:     f(module, *args, **kwargs)
[rank1]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 565, in __init__
[rank1]:     [get_layer(layer_id) for layer_id in range(num_layers)])
[rank1]:     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 565, in <listcomp>
[rank1]:     [get_layer(layer_id) for layer_id in range(num_layers)])
[rank1]:      ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 532, in get_layer
[rank1]:     return BaseTransformerLayer(
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank1]:     f(module, *args, **kwargs)
[rank1]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 363, in __init__
self.query_key_value = ColumnParallelLinear(                                                                                                                                               [76/1988]
[rank1]:                            ^^^^^^^^^^^^^^^^^^^^^
[rank1]: TypeError: ColumnParallelLinear.__init__() got multiple values for argument 'module'
[rank3]: Traceback (most recent call last):
[rank3]:   File "/mnt/afs//Concat-ID/train_video.py", line 235, in <module>
[rank3]:     training_main(
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/training/deepspeed_training.py", line 84, in training_main
[rank3]:     model = get_model(args, model_cls)
[rank3]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/base_model.py", line 418, in get_model
[rank3]:     model = model_cls(args, params_dtype=params_dtype, **kwargs)
[rank3]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank3]:     f(module, *args, **kwargs)
[rank3]:   File "/mnt/afs//Concat-ID/diffusion_video.py", line 74, in __init__
[rank3]:     model = instantiate_from_config(network_config)
[rank3]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs//Concat-ID/sgm/util.py", line 261, in instantiate_from_config
[rank3]:     return get_obj_from_str(config["target"])(**config.get("params", dict()), **extra_kwargs)
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank3]:     f(module, *args, **kwargs)
[rank3]:   File "/mnt/afs//Concat-ID/dit_video_concat.py", line 662, in __init__
[rank3]:     super().__init__(args=transformer_args, transformer=None, **kwargs)
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank3]:     f(module, *args, **kwargs)
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/base_model.py", line 93, in __init__
[rank3]:     self.transformer = BaseTransformer(
[rank3]:                        ^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank3]:     f(module, *args, **kwargs)
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 565, in __init__
[rank3]:     [get_layer(layer_id) for layer_id in range(num_layers)])
[rank3]:     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 565, in <listcomp>
[rank3]:     [get_layer(layer_id) for layer_id in range(num_layers)])
[rank3]:      ^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 532, in get_layer
[rank3]:     return BaseTransformerLayer(
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank3]:     f(module, *args, **kwargs)
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 363, in __init__
[rank3]:     self.attention = SelfAttention(
[rank3]:                      ^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 511, in wrapper
[rank3]:     f(module, *args, **kwargs)
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 66, in __init__

[rank3]:     self.query_key_value = ColumnParallelLinear(
[rank3]:                            ^^^^^^^^^^^^^^^^^^^^^
[rank3]: TypeError: ColumnParallelLinear.__init__() got multiple values for argument 'module'
[rank0]:[W412 09:52:22.237174988 ProcessGroupNCCL.cpp:1250] Warning: WARNING: process group has NOT been destroyed before we destruct ProcessGroupNCCL. On normal program exit, the application should call destr
oy_process_group to ensure that any pending NCCL operations have finished in this process. In rare cases this process can exit before this point and block the progress of another member of the process group. T
his constraint has always been present,  but this warning has only been added since PyTorch 2.4 (function operator())
pt-cd2d5528f04c4984a06d402621e9a7e6-worker-0:216524:216741 [0] NCCL INFO [Service thread] Connection closed by localRank 0
pt-cd2d5528f04c4984a06d402621e9a7e6-worker-0:216526:216737 [2] NCCL INFO [Service thread] Connection closed by localRank 2

[sunshinejuanjuan123]

修改video_size 为[240,360],报错

                                                                                                                                                                      
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/einops/einops.py", line 523, in reduce                                                                                            
[rank3]:     return _apply_recipe(                                                                                                                                                                               
[rank3]:            ^^^^^^^^^^^^^^                                                                                                                                                                               
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/einops/einops.py", line 234, in _apply_recipe                                                                                     
[rank3]:     init_shapes, axes_reordering, reduced_axes, added_axes, final_shapes, n_axes_w_added = _reconstruct_from_shape(                                                                                     
[rank3]:                                                                                            ^^^^^^^^^^^^^^^^^^^^^^^^                                                                                     
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/einops/einops.py", line 187, in _reconstruct_from_shape_uncached                                                                  
[rank3]:     raise EinopsError(f"Shape mismatch, can't divide axis of length {length} in chunks of {known_product}")                                                                                             
[rank3]: einops.EinopsError: Shape mismatch, can't divide axis of length 4620 in chunks of 1350                                                                                                                  
                                                                                                                                                                                                                 
[rank3]: During handling of the above exception, another exception occurred:                                                                                                                                     
                                                                                                                                                                                                                 
[rank3]: Traceback (most recent call last):                                                                                                                                                                      
[rank3]:   File "/mnt/afs//Concat-ID/train_video.py", line 235, in <module>                                                                                                                          
[rank3]:     training_main(                                                                                                                                                                                      
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/training/deepspeed_training.py", line 158, in training_main                                                                   
[rank3]:     iteration, skipped = train(model, optimizer,                                                                                                                                                        
[rank3]:                          ^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                        
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/training/deepspeed_training.py", line 360, in train                                                                           
[rank3]:     lm_loss, skipped_iter, metrics = train_step(train_data_iterator,                                                                                                                                    
[rank3]:                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                    
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/training/deepspeed_training.py", line 444, in train_step
[rank3]:     forward_ret = forward_step(data_iterator, model, args, timers, **kwargs)
[rank3]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs//Concat-ID/train_video.py", line 205, in forward_step
[rank3]:     loss, loss_dict = model.shared_step(batch)
[rank3]:                       ^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs//Concat-ID/diffusion_video.py", line 185, in shared_step
[rank3]:     loss, loss_dict = self(x, batch)
[rank3]:                       ^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
[rank3]:     return self._call_impl(*args, **kwargs) 
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
[rank3]:     return forward_call(*args, **kwargs)
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs//Concat-ID/diffusion_video.py", line 132, in forward
[rank3]:     loss = self.loss_fn(self.model, self.denoiser, self.conditioner, x, batch)
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs//Concat-ID/sgm/modules/diffusionmodules/loss.py", line 109, in __call__
[rank3]:     model_output = denoiser(network, noised_input, alphas_cumprod_sqrt, cond, **additional_model_inputs)
[rank3]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
[rank3]:     return forward_call(*args, **kwargs)
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/sat/model/transformer.py", line 738, in forward
[rank3]:     logits_parallel = self.hooks['final_forward'](logits, **kw_args, parallel_output=self.parallel_output)
[rank3]:                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs//Concat-ID/dit_video_concat.py", line 401, in final_forward
[rank3]:     return unpatchify(
[rank3]:            ^^^^^^^^^^^
[rank3]:   File "/mnt/afs//Concat-ID/dit_video_concat.py", line 367, in unpatchify
[rank3]:     imgs = rearrange(x, "b (t h w) (c p q) -> b t c (h p) (w q)", b=b, h=h, w=w, c=c, p=p, q=p) 
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/einops/einops.py", line 591, in rearrange
[rank3]:     return reduce(tensor, pattern, reduction="rearrange", **axes_lengths)
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/einops/einops.py", line 533, in reduce
[rank3]:     raise EinopsError(message + "\n {}".format(e))
[rank3]: einops.EinopsError:  Error while processing rearrange-reduction pattern "b (t h w) (c p q) -> b t c (h p) (w q)".
[rank3]:  Input tensor shape: torch.Size([1, 4620, 64]). Additional info: {'b': 1, 'h': 30, 'w': 45, 'c': 16, 'p': 2, 'q': 2}.
[rank3]:  Shape mismatch, can't divide axis of length 4620 in chunks of 1350
pt-cd2d5528f04c4984a06d402621e9a7e6-worker-0:218249:219061 [3] NCCL INFO [Service thread] Connection closed by localRank 3
[rank2]: Traceback (most recent call last):
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/einops/einops.py", line 523, in reduce
[rank2]:     return _apply_recipe(
[rank2]:            ^^^^^^^^^^^^^^
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/einops/einops.py", line 234, in _apply_recipe
[rank2]:     init_shapes, axes_reordering, reduced_axes, added_axes, final_shapes, n_axes_w_added = _reconstruct_from_shape(
[rank2]:                                                                                            ^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/mnt/afs/miniconda/envs/concatid/lib/python3.11/site-packages/einops/einops.py", line 187, in _reconstruct_from_shape_uncached
[rank2]:     raise EinopsError(f"Shape mismatch, can't divide axis of length {length} in chunks of {known_product}")
[rank2]: einops.EinopsError: Shape mismatch, can't divide axis of length 4620 in chunks of 1350

[rank2]: During handling of the above exception, another exception occurred:

[yongzhong-ruc]

I have not tested training Concat-ID with DeepSpeed of stage 3. However, we can successfully train Concat-ID with DeepSpeed of stage 2 on 8 H800 GPUs. The OutOfMemoryError seems to be caused by the VAE during inference in the training phase. Instead of performing inference during training, you can avoid it by setting eval_interval to a larger value (such as 100,000).