refactor rotary embedding 3: so it is not on cpu

[yiyixuxu]

fix#9299

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

Thanks!

@yiyixuxu possible to trigger a torch.compile() with PyTorch nightly to verify if this helps with the CUDAGraph issue? Code is in #9299 (comment).

Ccing @cpuhrsch maybe you would like to review it?

[yiyixuxu]

@sayakpaul yeah tested it is fine

[yiyixuxu]

@sayakpaul
Is this a reasonable script? I want to compare the performance against 0.30.1-patch before we introduce the rotary embedding refractor

importtorchimporttorch.utils.benchmarkasbenchmarkimportgcimporttimetorch.set_float32_matmul_precision("high")
torch._inductor.conv_1x1_as_mm=Truetorch._inductor.coordinate_descent_tuning=Truetorch._inductor.epilogue_fusion=Falsetorch._inductor.coordinate_descent_check_all_directions=Trueimportdiffusersfromplatformimportpython_versionfromdiffusersimportDiffusionPipelineprint(diffusers.__version__)
print(torch.__version__)
print(python_version())
defbenchmark_fn(f, *args, **kwargs):
t0=benchmark.Timer(
stmt="f(*args, **kwargs)",
globals={"args": args, "kwargs": kwargs, "f": f},
num_threads=torch.get_num_threads(),
)
returnf"{(t0.blocked_autorange().mean):.3f}"defbytes_to_giga_bytes(bytes):
returnf"{(bytes/1024/1024/1024):.3f}"defflush():
"""Wipes off memory."""gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_peak_memory_stats()
pipe=DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", text_encoder=None, text_encoder_2=None, torch_dtype=torch.bfloat16).to("cuda")
pipe.transformer.to(memory_format=torch.channels_last)
pipe.vae.to(memory_format=torch.channels_last)
pipe.transformer=torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True)
pipe.vae.decode=torch.compile(pipe.vae.decode, mode="max-autotune", fullgraph=True)
prompt_embeds=torch.load("flux_prompt_embeds.pt")
pooled_prompt_embeds=torch.load("flux_pooled_prompt_embeds.pt")
defrun_inference(pipe):
_=pipe(
prompt_embeds=prompt_embeds,
pooled_prompt_embeds=pooled_prompt_embeds,
num_inference_steps=5,
guidance_scale=3.5,
max_sequence_length=512,
generator=torch.manual_seed(42),
height=1024,
width=1024,
)
flush()
time=benchmark_fn(run_inference)
memory=bytes_to_giga_bytes(torch.cuda.max_memory_allocated()) # in GBs.print(f" Execution time: {time} sec")
print(f" Memory: {memory} gib")

[cpuhrsch]

I'd expect this to cause a sync as well since by default arange allocates on the CPU. One way to mitigate could be to
a) use pin_memory() on freqs ahead of time and set non_blocking=True
b) do arange on the GPU right away (i.e. torch.arange([...], device=pos.device)).

[yiyixuxu]

ohhh let's do torch.arange([...], device=pos.device)

[sayakpaul]

@yiyixuxu that looks reasonable but I'd call run_inference() maybe 2/3 times for warmups.

[cpuhrsch]

This should also be passed a device argument to allocate it on the GPU. If this isn't on the GPU, then neither will the following Tensors.