train_vae显存不够

[shushanxingzhe]

@bubbliiiing 我使用下面的命令训练vae，我看到autoencoder_kl_32x32x4_cogvideox.yaml的帧数video_len=49，但是我用48G显存的L40s GPU最大帧数只能开到13，我的视频分辨率是1440*1080，请问你们训练为什么帧数可以开到49？是视频分辨率更低吗？烦请拨冗回答一下

torchrun --master_addr=127.0.0.1 --master_port=6006 --nnodes=1 --node_rank=0 --nproc_per_node=4 scripts/train_vae.py
-t
-b easyanimate/vae/configs/autoencoder/autoencoder_kl_32x32x4_cogvideox.yaml
--scale_lr=False
--precision=bf16
--max_steps=2000
--resume="True"

[bubbliiiing]

你可以提供一下easyanimate/vae/configs/autoencoder/autoencoder_kl_32x32x4_cogvideox.yaml的细节看看吗？

[shushanxingzhe]

`model:
base_learning_rate: 1.0e-04
target: easyanimate.vae.ldm.models.cogvideox_casual3dcnn.AutoencoderKLMagvit_CogVideoX
params:
latent_channels: 16
temporal_compression_ratio: 4
monitor: train/rec_loss
ckpt_path: /data/aigc/models/智普/CogVideoX1.5-5B/vae/diffusion_pytorch_model.safetensors
down_block_types: ("CogVideoXDownBlock3D", "CogVideoXDownBlock3D", "CogVideoXDownBlock3D",
"CogVideoXDownBlock3D",)
up_block_types: ("CogVideoXUpBlock3D", "CogVideoXUpBlock3D", "CogVideoXUpBlock3D",
"CogVideoXUpBlock3D",)
lossconfig:
target: easyanimate.vae.ldm.modules.losses.LPIPSWithDiscriminator
params:
disc_start: 50001
kl_weight: 1.0e-06
disc_weight: 0.5
l2_loss_weight: 0.1
l1_loss_weight: 1.0
perceptual_weight: 1.0

data:
target: train_vae.DataModuleFromConfig

params:
batch_size: 1
wrap: true
num_workers: 4
train:
target: easyanimate.vae.ldm.data.dataset_image_video.CustomSRTrain
params:
data_json_path: json_of_internal_datasets.json
size: 256
degradation: pil_nearest
video_size: 256
video_len: 13
slice_interval: 1
validation:
target: easyanimate.vae.ldm.data.dataset_image_video.CustomSRValidation
params:
data_json_path: val_100.json
size: 256
degradation: pil_nearest
video_size: 256
video_len: 13
slice_interval: 1

lightning:
callbacks:
image_logger:
target: train_vae.ImageLogger
params:
batch_frequency: 5000
max_images: 8
increase_log_steps: True

trainer:
benchmark: True
accumulate_grad_batches: 1
gpus: "0,1,2,3"
num_nodes: 1`

[shushanxingzhe]

视频的分辨率都是14401080和10801440的

[bubbliiiing]

还有两点显存优化：
1、一个是使用bf16训练，应该是在trainer设置精度。
2、一个是使用gradient checkpointing，这个需要在代码里开启一下。https://github.com/aigc-apps/EasyAnimate/blob/main/easyanimate/vae/ldm/models/omnigen_enc_dec.py#L151C14-L151C36

[shushanxingzhe]

1、我已经用了bf16，用之前帧长只能开到5，改成bf16后才可以开到13
2、我训的是cogvideox不是omnigen，我改了http://github.com/aigc-apps/EasyAnimate/blob/main/easyanimate/vae/ldm/models/cogvideox_enc_dec.py#L270 设置gradient_checkpointing=True，但是显存占用只减少了1G

所以我想知道一下你们训练视频用的分辨率是多少，我按照目前的占用情况估算，就算我换80G显存的A100也没法将帧数开到您配置你的默认值49

[bubbliiiing]

你这里面分辨率已经在config里面固定256了，我们训练也是256的。
80G是可以训练49帧的，我没试过用L40s，没有这个卡，刚查了一下是48G，48G训练这个应该压力比较大。