1st training step OOMs after load_state: deepspeed stage 2 cpu offload
problem:1st training step OOMs after load_state: deepspeed stage 2 cpu offloadgpu memory usage significantly higher after load stateupdate:tried:model, train_dataloader, optimizer, lr_scheduler = accelerator.prepare(
model, train_dataloader, optimizer, lr_scheduler
)
model.load_checkpoint("path/to/input/dir/of/ckpt","pytorch_model")as suggested here:load_state cuda out of memory · Issue #1707 · huggingface/accelerate · GitHubstill OOMsDescription of problem:Code trains correctly, can use batch size up 22 without OOMSaves state without errorsLoads state without errors1st training step post load_state OOMs if batch size > 2Anyone know how to correctly load_state with stage 2 cpu offload?Relevant code:save_state code:#(post epoch)if save_state == True:
    accelerator.save_state(save_path)load_state code:if load_saved_state == True:
	#load base_model from base_model pipeline (same as initial train loading)
        noise_scheduler = DDPMScheduler.from_pretrained(
            pretrained_model_name_or_path, subfolder="scheduler"
        )
        unet = UNet2DConditionModel.from_pretrained(
            pretrained_model_name_or_path, subfolder="unet"
        )
        optimizer = bnb.optim.AdamW8bit(unet.parameters(), lr=learning_rate)
        lr_scheduler = get_scheduler(
                learning_rate_scheduler,
                optimizer=optimizer,
                num_warmup_steps=num_warmup_update_steps * num_processes,
                num_training_steps=total_num_update_steps,
            )
           
         #accelerate.prepare
         unet, optimizer, train_dataloader, validation_loss_dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, train_dataloader, validation_loss_dataloader, lr_scheduler)
         
         #load_state
         accelerator.load_state(load_saved_state)accelerate config:compute_environment: LOCAL_MACHINEdebug: truedeepspeed_config:offload_optimizer_device: cpuoffload_param_device: cpuzero3_init_flag: falsezero_stage: 2distributed_type: DEEPSPEEDdowncast_bf16: ‘no’enable_cpu_affinity: falsemachine_rank: 0main_training_function: mainmixed_precision: fp16num_machines: 1num_processes: 1rdzv_backend: staticsame_network: truetpu_env:tpu_use_cluster: falsetpu_use_sudo: falseuse_cpu: false