训练时OOM

[tonney007]

遇到了oom，我应该怎么做呢？

训练命令

torchrun --standalone --nproc_per_node 2 scripts/train.py configs/opensora-v1-1/train/stage1.py --data-path /opt/ml/code/Open-Sora/data/meta/meta_clips_caption_cleaned.csv --ckpt-path /opt/ml/code/OpenSora-STDiT-v2-stage3

硬件用的 2*A800 （80G）

config：

~# Define dataset dataset = dict( type="VariableVideoTextDataset", data_path=None, num_frames=None, frame_interval=3, image_size=(None, None), transform_name="resize_crop", ) ~# IMG: 1024 (20%) 512 (30%) 256 (50%) drop (50%) bucket_config = { # 1s/it "144p": {1: (0.5, 48), 16: (1.0, 6), 32: (1.0, 3), 96: (1.0, 1)}, "256": {1: (0.5, 24), 16: (0.5, 3), 48: (0.5, 1), 64: (0.0, None)}, "240p": {16: (0.3, 2), 32: (0.3, 1), 64: (0.0, None)}, "512": {1: (0.4, 12)}, "1024": {1: (0.3, 3)}, } mask_ratios = { "mask_no": 0.75, "mask_quarter_random": 0.025, "mask_quarter_head": 0.025, "mask_quarter_tail": 0.025, "mask_quarter_head_tail": 0.05, "mask_image_random": 0.025, "mask_image_head": 0.025, "mask_image_tail": 0.025, "mask_image_head_tail": 0.05, }

~# Define acceleration num_workers = 2 num_bucket_build_workers = 4 ~# num_workers = 8 ~# num_bucket_build_workers = 16 dtype = "bf16" grad_checkpoint = False plugin = "zero2" sp_size = 1

~# Define model model = dict( type="STDiT2-XL/2", from_pretrained=None, input_sq_size=256, # pretrained model is trained on 512x512 qk_norm=True, enable_flash_attn=False, enable_layernorm_kernel=False , ) vae = dict( type="VideoAutoencoderKL", from_pretrained="/opt/ml/code/sd-vae-ft-ema", micro_batch_size=1, ~# micro_batch_size=4, local_files_only=True, ) text_encoder = dict( type="t5", from_pretrained="/opt/ml/code/t5-v1_1-xxl", model_max_length=200, shardformer=False, local_files_only=True, ) scheduler = dict( type="iddpm", timestep_respacing="", )

~# Others seed = 42 outputs = "outputs" wandb = False

epochs = 1000 log_every = 10 ckpt_every = 500 load = None

batch_size = 1 lr = 2e-5 grad_clip = 1.0

torch.cuda.OutOfMemoryError: CUDA out of memory.

[tonney007]

可以了，需要将grad_checkpoint = False打开，显存会骤降一半以上

[TYang92677626]

grad_checkpoint = False

请问这个参数是什么意思呢

[TYang92677626]

推理用时多久呢

[tonney007]

推理用时多久呢

推理时长跟显卡、视频长度、分辨率有关系，我是a800，十几秒到十几分钟不等。

[tonney007]

grad_checkpoint

gpt回~~使用梯度检查点来减少内存使用。

[aulaywang]

有没有支持更低显存（比如只有45G）训练的修改方式呢？现在grad_checkpoint也加了，num_workers = 0，num_bucket_build_workers = 1