for epoch in range(num_train_epochs):

        unet.train()

        for step, batch in enumerate(train_dataloader):

            with accelerator.accumulate(unet):

                # Convert images to latent space

                with torch.no_grad():

                    latents = vae.encode(batch["pixel_values"]).latent_dist.sample()

                    latents = latents * 0.18215

                # Sample noise that we'll add to the latents

                noise = torch.randn(latents.shape).to(latents.device)

                bsz = latents.shape[0]

                # Sample a random timestep for each image

                timesteps = torch.randint(

                    0,

                    noise_scheduler.config.num_train_timesteps,

                    (bsz,),

                    device=latents.device,

                ).long()

                # Add noise to the latents according to the noise magnitude at each timestep

                # (this is the forward diffusion process)

                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)

                # Get the text embedding for conditioning

                with torch.no_grad():

                    encoder_hidden_states = text_encoder(batch["input_ids"])[0]

                # Predict the noise residual

                noise_pred = unet(

                    noisy_latents, timesteps, encoder_hidden_states

                ).sample

                loss = (

                    F.mse_loss(noise_pred, noise, reduction="none")

                    .mean([1, 2, 3])

                    .mean()

                )

                accelerator.backward(loss)

                # if accelerator.sync_gradients:

                #     accelerator.clip_grad_norm_(unet.parameters(), args.max_grad_norm)

                optimizer.step()

                optimizer.zero_grad()

            # Checks if the accelerator has performed an optimization step behind the scenes

            if accelerator.sync_gradients:

                progress_bar.update(1)

                global_step += 1

            logs = {"loss": loss.detach().item()}

            progress_bar.set_postfix(**logs)

            if global_step >= args.max_train_steps:

                break

        accelerator.wait_for_everyone()

    # Create the pipeline using the trained modules and save it

    if accelerator.is_main_process:

        print(f"Loading pipeline and saving to {args.output_dir}...")

        scheduler = PNDMScheduler(

            beta_start=0.00085,

            beta_end=0.012,

            beta_schedule="scaled_linear",

            skip_prk_steps=True,

            steps_offset=1,

        )

        pipeline = StableDiffusionPipeline(

            text_encoder=text_encoder,

            vae=vae,

            unet=accelerator.unwrap_model(unet),

            tokenizer=tokenizer,

            scheduler=scheduler,

            safety_checker=StableDiffusionSafetyChecker.from_pretrained(

                "CompVis/stable-diffusion-safety-checker"

            ),

            feature_extractor=feature_extractor,

        )

        pipeline.save_pretrained(args.output_dir)

调试过程中执行到accelerator.backward(loss)就会报错

[E compiler_depend.ts:280] call aclnnGeGluV3 failed, detail:EZ9999: Inner Error! 

EZ9999: 2024-11-05-18:43:18.297.398  Op GeGluV2 does not has any binary. 

        TraceBack (most recent call last): 

        Kernel Run failed. opType: 47, GeGluV2 

        launch failed for GeGluV2, errno:561000. 

 

[ERROR] 2024-11-05-18:43:18 (PID:48061, Device:0, RankID:-1) ERR01100 OPS call acl api failed 

Exception raised from operator() at build/CMakeFiles/torch_npu.dir/compiler_depend.ts:39 (most recent call first): 

frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x68 (0xfffe26510898 in /root/anaconda3/lib/python3.9/site-packages/torch/lib/libc10.so) 

frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) + 0x6c (0xfffe264c92a8 in /root/anaconda3/lib/python3.9/site-packages/torch/lib/libc10.so) 

frame #2: <unknown function> + 0xc8f94c (0xfffce0c9394c in /root/anaconda3/lib/python3.9/site-packages/torch_npu/lib/libtorch_npu.so) 

frame #3: <unknown function> + 0x136ed1c (0xfffce1372d1c in /root/anaconda3/lib/python3.9/site-packages/torch_npu/lib/libtorch_npu.so) 

frame #4: <unknown function> + 0x631b54 (0xfffce0635b54 in /root/anaconda3/lib/python3.9/site-packages/torch_npu/lib/libtorch_npu.so) 

frame #5: <unknown function> + 0x6320bc (0xfffce06360bc in /root/anaconda3/lib/python3.9/site-packages/torch_npu/lib/libtorch_npu.so) 

frame #6: <unknown function> + 0x62f534 (0xfffce0633534 in /root/anaconda3/lib/python3.9/site-packages/torch_npu/lib/libtorch_npu.so) 

frame #7: <unknown function> + 0x946ec (0xfffe265376ec in /root/anaconda3/lib/python3.9/site-packages/torch/lib/libc10.so) 

frame #8: <unknown function> + 0x7a80 (0xffff9492aa80 in /lib64/libpthread.so.0) 

frame #9: <unknown function> + 0xe4d0c (0xffff9475ed0c in /lib64/libc.so.6)