Context Parallel w/ Ring & Ulysses & Unified Attention

[a-r-r-o-w]

Adds support for ring, ulysses and unified attention natively. For a minimal PoC, I've limited changes to Flux.

Supported attention backends with CP: cuDNN, FA2, Sage.

Requires #11916 to be merged first.

Minimal example

Note: the examples here are not up-to-date! Please refer to the official examples once the docs are uodated

importtorchfromdiffusersimportFluxPipelinefromdiffusersimportParallelConfig, enable_parallelismtry:
torch.distributed.init_process_group("nccl")
rank=torch.distributed.get_rank()
device=torch.device("cuda", rank%torch.cuda.device_count())
torch.cuda.set_device(device)
pipe=FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.to(device)
pipe.transformer.parallelize(config=ParallelConfig(ulysses_degree=2))
# pipe.transformer.set_attention_backend("_native_cudnn")pipe.transformer.set_attention_backend("flash")
# pipe.transformer.compile(fullgraph=True)# pipe.transformer.compile(fullgraph=True, mode="max-autotune")prompt="A cat holding a sign that says 'hello world'"# Must specify generator so all ranks start with same latents (or pass your own)generator=torch.Generator().manual_seed(42)
withenable_parallelism(pipe):
image=pipe(prompt, num_inference_steps=2, guidance_scale=4.0, generator=generator).images[0]
image=pipe(prompt, num_inference_steps=28, guidance_scale=4.0, generator=generator).images[0]
ifrank==0:
image.save("output.png")
exceptExceptionase:
print(f"An error occurred: {e}")
torch.distributed.breakpoint()
raisefinally:
iftorch.distributed.is_initialized():
torch.distributed.destroy_process_group()

Wan

importtorchfromdiffusersimportAutoencoderKLWan, WanPipeline, ParallelConfig, enable_parallelismfromdiffusers.utilsimportexport_to_videofromdiffusers.schedulers.scheduling_unipc_multistepimportUniPCMultistepSchedulertry:
torch.distributed.init_process_group("nccl")
rank=torch.distributed.get_rank()
device=torch.device("cuda", rank%torch.cuda.device_count())
torch.cuda.set_device(device)
# Available models: Wan-AI/Wan2.1-T2V-14B-Diffusers, Wan-AI/Wan2.1-T2V-1.3B-Diffusersmodel_id="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"vae=AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
pipe=WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
flow_shift=5.0# 5.0 for 720P, 3.0 for 480Ppipe.scheduler=UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
pipe.to("cuda")
pipe.transformer.parallelize(config=ParallelConfig(ulysses_degree=2))
pipe.transformer.set_attention_backend("flash")
prompt="A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window."negative_prompt="Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"# Must specify generator so all ranks start with same latents (or pass your own)generator=torch.Generator().manual_seed(42)
withenable_parallelism(pipe):
output=pipe(
prompt=prompt,
negative_prompt=negative_prompt,
height=480,
width=832,
num_frames=81,
guidance_scale=5.0,
generator=generator,
).frames[0]
ifrank==0:
export_to_video(output, "output.mp4", fps=16)
exceptExceptionase:
print(f"An error occurred: {e}")
torch.distributed.breakpoint()
raisefinally:
iftorch.distributed.is_initialized():
torch.distributed.destroy_process_group()

Qwen

importtorchfromdiffusersimportQwenImagePipeline, ParallelConfig, enable_parallelismtry:
torch.distributed.init_process_group("nccl")
rank=torch.distributed.get_rank()
device=torch.device("cuda", rank%torch.cuda.device_count())
torch.cuda.set_device(device)
pipe=QwenImagePipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=torch.bfloat16)
pipe.to("cuda")
pipe.transformer.parallelize(config=ParallelConfig(ulysses_degree=2))
pipe.transformer.set_attention_backend("flash")
prompt="A cat holding a sign that says 'hello world'"# Must specify generator so all ranks start with same latents (or pass your own)generator=torch.Generator().manual_seed(42)
withenable_parallelism(pipe):
image=pipe(prompt, num_inference_steps=50, generator=generator).images[0]
ifrank==0:
image.save("output.png")
exceptExceptionase:
print(f"An error occurred: {e}")
torch.distributed.breakpoint()
raisefinally:
iftorch.distributed.is_initialized():
torch.distributed.destroy_process_group()

LTXVideo

importtorchfromdiffusersimportLTXPipeline, ParallelConfig, enable_parallelismfromdiffusers.utilsimportexport_to_videotry:
torch.distributed.init_process_group("nccl")
rank=torch.distributed.get_rank()
device=torch.device("cuda", rank%torch.cuda.device_count())
torch.cuda.set_device(device)
pipe=LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)
pipe.to("cuda")
pipe.transformer.parallelize(config=ParallelConfig(ulysses_degree=4))
pipe.transformer.set_attention_backend("flash")
prompt="A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage"negative_prompt="worst quality, inconsistent motion, blurry, jittery, distorted"# Must specify generator so all ranks start with same latents (or pass your own)generator=torch.Generator().manual_seed(42)
withenable_parallelism(pipe):
video=pipe(
prompt=prompt,
negative_prompt=negative_prompt,
width=704,
height=512,
num_frames=161,
num_inference_steps=50,
generator=generator,
).frames[0]
ifrank==0:
export_to_video(video, "output.mp4", fps=24)
exceptExceptionase:
print(f"An error occurred: {e}")
torch.distributed.breakpoint()
raisefinally:
iftorch.distributed.is_initialized():
torch.distributed.destroy_process_group()

Benchmarks

Flux code

importargparseimportcontextlibimportfunctoolsimportpathlibimportmathfromdataclassesimportdataclassfromtypingimportCallable, List, Literal, Optional, Tupleimportnumpyasnpimporttorchimporttorch.distributedasdistimporttorch.distributed._functional_collectivesasfuncolimporttorch.profiler._utilsimporttorch._dynamo.configimporttorch._inductor.configimporttorch._higher_order_ops.auto_functionalizeasaffromtorch.profilerimportprofile, record_function, ProfilerActivityfromdiffusersimportAutoencoderKLfromdiffusers.image_processorimportVaeImageProcessorfromdiffusers.configuration_utilsimportConfigMixin, register_to_configfromdiffusers.loadersimportFluxTransformer2DLoadersMixin, FromOriginalModelMixin, PeftAdapterMixinfromdiffusers.models.attentionimportFeedForwardfromdiffusers.models.embeddingsimportget_1d_rotary_pos_embedfromdiffusers.models.cache_utilsimportCacheMixinfromdiffusers.models.embeddingsimport (
CombinedTimestepGuidanceTextProjEmbeddings,
CombinedTimestepTextProjEmbeddings,
)
fromdiffusers.models.modeling_utilsimportModelMixinfromkernelsimportget_kernelfromtransformersimportCLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFasttry:
fromflash_attnimportflash_attn_funcexcept:
print("Flash Attention 2 not found.")
try:
fromflash_attn_interfaceimportflash_attn_funcasflash_attn_3_funcexcept:
print("Flash Attention 3 not found.")
defapply_flags():
torch._dynamo.config.inline_inbuilt_nn_modules=Falsetorch._dynamo.config.cache_size_limit=128torch._dynamo.config.error_on_recompile=Truetorch._inductor.config.conv_1x1_as_mm=Truetorch._inductor.config.coordinate_descent_check_all_directions=Truetorch._inductor.config.coordinate_descent_tuning=Truetorch._inductor.config.disable_progress=Falsetorch._inductor.config.fx_graph_cache=Truetorch._inductor.config.epilogue_fusion=Falsetorch._inductor.config.aggressive_fusion=Truetorch._inductor.config.shape_padding=Truetorch._inductor.config.triton.enable_persistent_tma_matmul=Truetorch.backends.cuda.matmul.allow_tf32=Truetorch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction=Truetorch.backends.cudnn.allow_tf32=Trueaf.auto_functionalized_v2._cacheable=Trueaf.auto_functionalized._cacheable=TrueROPE_PRECISION=torch.bfloat16SPATIAL_COMPRESSION_RATIO=8PIXEL_UNSHUFFLING_DOWNSAMPLING_FACTOR=2T5_SEQUENCE_LENGTH=512DEFAULT_HEIGHT=1024DEFAULT_WIDTH=1024PATCH_SIZE=1LATENT_HEIGHT=DEFAULT_HEIGHT// (SPATIAL_COMPRESSION_RATIO*PATCH_SIZE) //2LATENT_WIDTH=DEFAULT_WIDTH// (SPATIAL_COMPRESSION_RATIO*PATCH_SIZE) //2SUPPORTED_BUCKET_LENGTHS=list(range(128, 512+1, 64))
SUPPORTED_GUIDANCE_SCALES= [i/2foriinrange(41)] # 0, 0.5, 1.0, ..., 20.0MIN_INFERENCE_STEPS=2MAX_INFERENCE_STEPS=50ATTENTION_OP=NoneBASE_IMAGE_SEQ_LEN=256MAX_IMAGE_SEQ_LEN=4096BASE_SHIFT=0.5MAX_SHIFT=1.15M= (MAX_SHIFT-BASE_SHIFT) / (MAX_IMAGE_SEQ_LEN-BASE_IMAGE_SEQ_LEN)
B=BASE_SHIFT-M*BASE_IMAGE_SEQ_LEN@dataclassclassContextParallelOptions:
ring_degree: int=Noneulysses_degree: int=Nonemode: Literal["ring", "ulysses", "unified"] ="ring"mesh: dist.DeviceMesh|None=Noneconvert_to_fp32: bool=Trueattention_op: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]] |None=None_flattened_mesh: dist.DeviceMesh=None_ring_mesh: dist.DeviceMesh=None_ulysses_mesh: dist.DeviceMesh=None_ring_local_rank: int=None_ulysses_local_rank: int=Nonecp_options=ContextParallelOptions()
classAdaLayerNormContinuous(torch.nn.Module):
def__init__(
self, embedding_dim: int, conditioning_embedding_dim: int, elementwise_affine=True, eps=1e-5, bias=True
):
super().__init__()
self.linear=torch.nn.Linear(conditioning_embedding_dim, embedding_dim*2, bias=bias)
self.norm=torch.nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
defforward(self, x: torch.Tensor, emb: torch.Tensor) ->torch.Tensor:
emb=self.linear(emb)
scale, shift=emb.unsqueeze(1).chunk(2, dim=-1)
x=self.norm(x)
x=torch.addcmul(shift, x, 1+scale)
returnxclassAdaLayerNormZeroSingle(torch.nn.Module):
def__init__(self, embedding_dim: int, bias=True):
super().__init__()
self.linear=torch.nn.Linear(embedding_dim, 3*embedding_dim, bias=bias)
self.norm=torch.nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
defforward(self, x: torch.Tensor, emb: torch.Tensor):
shift_msa, scale_msa, gate_msa=emb.chunk(3, dim=-1)
x=self.norm(x)
x=torch.addcmul(shift_msa, x, 1+scale_msa)
returnx, gate_msaclassAdaLayerNormZero(torch.nn.Module):
def__init__(self, embedding_dim: int, bias=True):
super().__init__()
self.linear=torch.nn.Linear(embedding_dim, 6*embedding_dim, bias=bias)
self.norm=torch.nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
defforward(self, x: torch.Tensor, emb: torch.Tensor):
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp=emb.chunk(6, dim=-1)
x=self.norm(x)
x=torch.addcmul(shift_msa, x, 1+scale_msa)
returnx, gate_msa, shift_mlp, scale_mlp, gate_mlpclassAttention(torch.nn.Module):
def__init__(
self,
query_dim: int,
heads: int=8,
dim_head: int=64,
dropout: float=0.0,
bias: bool=False,
qk_norm: Optional[str] =None,
added_kv_proj_dim: Optional[int] =None,
added_proj_bias: Optional[bool] =True,
out_bias: bool=True,
eps: float=1e-5,
out_dim: int=None,
context_pre_only=None,
pre_only=False,
elementwise_affine: bool=True,
):
super().__init__()
assertqk_norm=="rms_norm", "Flux uses RMSNorm"self.inner_dim=out_dimifout_dimisnotNoneelsedim_head*headsself.query_dim=query_dimself.use_bias=biasself.dropout=dropoutself.out_dim=out_dimifout_dimisnotNoneelsequery_dimself.context_pre_only=context_pre_onlyself.pre_only=pre_onlyself.heads=out_dim//dim_headifout_dimisnotNoneelseheadsself.added_proj_bias=added_proj_biasself.norm_q=torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
self.norm_k=torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
self.to_q=torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_k=torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_v=torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
ifnotself.pre_only:
self.to_out=torch.nn.ModuleList([])
self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
ifadded_kv_proj_dimisnotNone:
self.norm_added_q=torch.nn.RMSNorm(dim_head, eps=eps)
self.norm_added_k=torch.nn.RMSNorm(dim_head, eps=eps)
self.add_q_proj=torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
self.add_k_proj=torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
self.add_v_proj=torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
self.to_add_out=torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias)
defforward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: Optional[torch.Tensor] =None,
attention_mask: Optional[torch.Tensor] =None,
image_rotary_emb: Optional[torch.Tensor] =None,
) ->torch.Tensor:
cos, sin=image_rotary_embifimage_rotary_embisnotNoneelse (None, None)
ifself.fused_projections:
query, key, value=self.to_qkv(hidden_states).chunk(3, dim=-1)
else:
query=self.to_q(hidden_states)
key=self.to_k(hidden_states)
value=self.to_v(hidden_states)
query, key, value= (x.unflatten(2, (self.heads, -1)) forxin (query, key, value))
query=self.norm_q(query)
key=self.norm_k(key)
ifencoder_hidden_statesisnotNone:
ifself.fused_projections:
query_c, key_c, value_c=self.to_added_qkv(encoder_hidden_states).chunk(3, dim=-1)
else:
query_c=self.add_q_proj(encoder_hidden_states)
key_c=self.add_k_proj(encoder_hidden_states)
value_c=self.add_v_proj(encoder_hidden_states)
query_c, key_c, value_c= (
x.unflatten(2, (self.heads, -1)) forxin (query_c, key_c, value_c)
)
query_c=self.norm_added_q(query_c)
key_c=self.norm_added_k(key_c)
query=torch.cat([query_c, query], dim=1)
key=torch.cat([key_c, key], dim=1)
value=torch.cat([value_c, value], dim=1)
ifimage_rotary_embisnotNone:
x_real, x_imag=query.unflatten(-1, (-1, 2)).unbind(-1)
x_rotated=torch.stack([-x_imag, x_real], dim=-1).flatten(3)
query= (query.to(ROPE_PRECISION) *cos+x_rotated.to(ROPE_PRECISION) *sin).type_as(query)
x_real, x_imag=key.unflatten(-1, (-1, 2)).unbind(-1)
x_rotated=torch.stack([-x_imag, x_real], dim=-1).flatten(3)
key= (key.to(ROPE_PRECISION) *cos+x_rotated.to(ROPE_PRECISION) *sin).type_as(key)
hidden_states, lse=ATTENTION_OP(query, key, value)
hidden_states=hidden_states.flatten(2)
ifencoder_hidden_statesisnotNone:
encoder_hidden_states, hidden_states=torch.split_with_sizes(
hidden_states,
[encoder_hidden_states.shape[1], hidden_states.shape[1] -encoder_hidden_states.shape[1]],
dim=1,
)
hidden_states=self.to_out[0](hidden_states)
encoder_hidden_states=self.to_add_out(encoder_hidden_states)
returnhidden_states, encoder_hidden_statesreturnhidden_states@torch.no_grad()deffuse_projections(self):
device=self.to_q.weight.data.devicedtype=self.to_q.weight.data.dtypeconcatenated_weights=torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data])
in_features=concatenated_weights.shape[1]
out_features=concatenated_weights.shape[0]
self.to_qkv=torch.nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype)
self.to_qkv.weight.copy_(concatenated_weights)
ifself.use_bias:
concatenated_bias=torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data])
self.to_qkv.bias.copy_(concatenated_bias)
if (
getattr(self, "add_q_proj", None) isnotNoneandgetattr(self, "add_k_proj", None) isnotNoneandgetattr(self, "add_v_proj", None) isnotNone
):
concatenated_weights=torch.cat(
[self.add_q_proj.weight.data, self.add_k_proj.weight.data, self.add_v_proj.weight.data]
)
in_features=concatenated_weights.shape[1]
out_features=concatenated_weights.shape[0]
self.to_added_qkv=torch.nn.Linear(
in_features, out_features, bias=self.added_proj_bias, device=device, dtype=dtype
)
self.to_added_qkv.weight.copy_(concatenated_weights)
ifself.added_proj_bias:
concatenated_bias=torch.cat(
[self.add_q_proj.bias.data, self.add_k_proj.bias.data, self.add_v_proj.bias.data]
)
self.to_added_qkv.bias.copy_(concatenated_bias)
forlayerin ("to_q", "to_k", "to_v", "to_added_q", "to_added_k", "to_added_v"):
ifhasattr(self, layer):
module=getattr(self, layer)
module.to("meta")
delattr(self, layer)
self.fused_projections=TrueclassFluxPosEmbed(torch.nn.Module):
def__init__(self, theta: int, axes_dim: List[int]):
super().__init__()
self.theta=thetaself.axes_dim=axes_dimdefforward(self, ids: torch.Tensor) ->torch.Tensor:
n_axes=ids.shape[-1]
cos_out= []
sin_out= []
foriinrange(n_axes):
cos, sin=get_1d_rotary_pos_embed(
self.axes_dim[i],
ids[:, i],
theta=self.theta,
repeat_interleave_real=True,
use_real=True,
freqs_dtype=ROPE_PRECISION,
)
cos_out.append(cos)
sin_out.append(sin)
freqs_cos=torch.cat(cos_out, dim=-1).to(ids.device, dtype=ROPE_PRECISION)[None, :, None]
freqs_sin=torch.cat(sin_out, dim=-1).to(ids.device, dtype=ROPE_PRECISION)[None, :, None]
returnfreqs_cos, freqs_sinclassFluxSingleTransformerBlock(torch.nn.Module):
def__init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float=4.0):
super().__init__()
self.mlp_hidden_dim=int(dim*mlp_ratio)
self.norm=AdaLayerNormZeroSingle(dim)
self.proj_mlp=torch.nn.Linear(dim, self.mlp_hidden_dim)
self.act_mlp=torch.nn.GELU(approximate="tanh")
self.attn=Attention(
query_dim=dim,
dim_head=attention_head_dim,
heads=num_attention_heads,
out_dim=dim,
bias=True,
qk_norm="rms_norm",
eps=1e-6,
pre_only=True,
)
self.proj_out=torch.nn.Linear(dim+self.mlp_hidden_dim, dim)
defforward(
self,
hidden_states: torch.Tensor,
temb: torch.Tensor,
image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] =None,
) ->torch.Tensor:
norm_hidden_states, gate=self.norm(hidden_states, emb=temb)
mlp_hidden_states=self.act_mlp(self.proj_mlp(norm_hidden_states))
attn_output=self.attn(
hidden_states=norm_hidden_states,
image_rotary_emb=image_rotary_emb,
)
attn_mlp_hidden_states=torch.cat([attn_output, mlp_hidden_states], dim=2)
proj_out=self.proj_out(attn_mlp_hidden_states)
hidden_states=torch.addcmul(hidden_states, gate, proj_out)
returnhidden_statesclassFluxTransformerBlock(torch.nn.Module):
def__init__(
self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str="rms_norm", eps: float=1e-6
):
super().__init__()
self.norm1=AdaLayerNormZero(dim)
self.norm1_context=AdaLayerNormZero(dim)
self.attn=Attention(
query_dim=dim,
added_kv_proj_dim=dim,
dim_head=attention_head_dim,
heads=num_attention_heads,
out_dim=dim,
context_pre_only=False,
bias=True,
qk_norm=qk_norm,
eps=eps,
)
self.norm2=torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
self.ff=FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
self.norm2_context=torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
self.ff_context=FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
defforward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] =None,
) ->Tuple[torch.Tensor, torch.Tensor]:
temb, temb_context=temb.chunk(2, dim=-1)
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp=self.norm1(hidden_states, emb=temb)
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp=self.norm1_context(
encoder_hidden_states, emb=temb_context
)
attn_output, context_attn_output=self.attn(
hidden_states=norm_hidden_states,
encoder_hidden_states=norm_encoder_hidden_states,
image_rotary_emb=image_rotary_emb,
)
hidden_states=torch.addcmul(hidden_states, gate_msa, attn_output)
encoder_hidden_states=torch.addcmul(encoder_hidden_states, c_gate_msa, context_attn_output)
norm_hidden_states=self.norm2(hidden_states)
norm_encoder_hidden_states=self.norm2_context(encoder_hidden_states)
norm_hidden_states=torch.addcmul(shift_mlp, norm_hidden_states, 1+scale_mlp)
norm_encoder_hidden_states=torch.addcmul(c_shift_mlp, norm_encoder_hidden_states, 1+c_scale_mlp)
ff_output=self.ff(norm_hidden_states)
context_ff_output=self.ff_context(norm_encoder_hidden_states)
hidden_states=torch.addcmul(hidden_states, gate_mlp, ff_output)
encoder_hidden_states=torch.addcmul(encoder_hidden_states, c_gate_mlp, context_ff_output)
returnencoder_hidden_states, hidden_statesclassFluxTransformer2DModel(
ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, FluxTransformer2DLoadersMixin, CacheMixin
):
@register_to_configdef__init__(
self,
patch_size: int=1,
in_channels: int=64,
out_channels: Optional[int] =None,
num_layers: int=19,
num_single_layers: int=38,
attention_head_dim: int=128,
num_attention_heads: int=24,
joint_attention_dim: int=4096,
pooled_projection_dim: int=768,
guidance_embeds: bool=False,
axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
):
super().__init__()
self.out_channels=out_channelsorin_channelsself.inner_dim=num_attention_heads*attention_head_dimself.pos_embed=FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
text_time_guidance_cls= (
CombinedTimestepGuidanceTextProjEmbeddingsifguidance_embedselseCombinedTimestepTextProjEmbeddings
)
self.time_text_embed=text_time_guidance_cls(
embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim
)
self.context_embedder=torch.nn.Linear(joint_attention_dim, self.inner_dim)
self.x_embedder=torch.nn.Linear(in_channels, self.inner_dim)
self.transformer_blocks=torch.nn.ModuleList(
[
FluxTransformerBlock(
dim=self.inner_dim,
num_attention_heads=num_attention_heads,
attention_head_dim=attention_head_dim,
)
for_inrange(num_layers)
]
)
self.single_transformer_blocks=torch.nn.ModuleList(
[
FluxSingleTransformerBlock(
dim=self.inner_dim,
num_attention_heads=num_attention_heads,
attention_head_dim=attention_head_dim,
)
for_inrange(num_single_layers)
]
)
self.norm_out=AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
self.proj_out=torch.nn.Linear(self.inner_dim, patch_size*patch_size*self.out_channels, bias=True)
defforward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
conditioning: torch.Tensor,
image_rotary_emb: Tuple[torch.Tensor, torch.Tensor],
dt: torch.Tensor,
) ->torch.Tensor:
x_t=hidden_stateshidden_states=self.x_embedder(hidden_states)
adaln_linear_dual_stream_states=self.adaln_linear(conditioning).unsqueeze(1).chunk(self.config.num_layers, dim=-1)
adaln_linear_single_stream_states=self.adaln_linear_single(conditioning).unsqueeze(1).chunk(self.config.num_single_layers, dim=-1)
fori, blockinenumerate(self.transformer_blocks):
encoder_hidden_states, hidden_states=block(
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
temb=adaln_linear_dual_stream_states[i],
image_rotary_emb=image_rotary_emb,
)
hidden_states=torch.cat([encoder_hidden_states, hidden_states], dim=1)
fori, blockinenumerate(self.single_transformer_blocks):
hidden_states=block(
hidden_states=hidden_states,
temb=adaln_linear_single_stream_states[i],
image_rotary_emb=image_rotary_emb,
)
hidden_states=hidden_states[:, encoder_hidden_states.shape[1] :, ...]
hidden_states=self.norm_out(hidden_states, conditioning)
velocity=self.proj_out(hidden_states)
x=x_t+dt*velocityreturnx@torch.no_grad()deffuse_qkv_(model: FluxTransformer2DModel) ->FluxTransformer2DModel:
forsubmoduleinmodel.modules():
ifnotisinstance(submodule, Attention):
continuesubmodule.fuse_projections()
@torch.no_grad()deffuse_adaln_linear_(model: FluxTransformer2DModel) ->FluxTransformer2DModel:
adaln_linear_weights= []
adaln_linear_biases= []
forblockinmodel.transformer_blocks:
adaln_linear_weights.append(block.norm1.linear.weight.data.clone())
adaln_linear_weights.append(block.norm1_context.linear.weight.data.clone())
adaln_linear_biases.append(block.norm1.linear.bias.data.clone())
adaln_linear_biases.append(block.norm1_context.linear.bias.data.clone())
block.norm1.linear.to("meta")
block.norm1_context.linear.to("meta")
delblock.norm1.linear, block.norm1_context.linearadaln_linear_weights=torch.cat(adaln_linear_weights, dim=0)
adaln_linear_biases=torch.cat(adaln_linear_biases, dim=0)
in_features=adaln_linear_weights.shape[1]
out_features=adaln_linear_weights.shape[0]
model.adaln_linear=torch.nn.Linear(
in_features, out_features, bias=True, device=adaln_linear_weights.device, dtype=adaln_linear_weights.dtype
)
model.adaln_linear.weight.copy_(adaln_linear_weights)
model.adaln_linear.bias.copy_(adaln_linear_biases)
adaln_linear_weights= []
adaln_linear_biases= []
forblockinmodel.single_transformer_blocks:
adaln_linear_weights.append(block.norm.linear.weight.data.clone())
adaln_linear_biases.append(block.norm.linear.bias.data.clone())
block.norm.linear.to("meta")
delblock.norm.linearadaln_linear_weights=torch.cat(adaln_linear_weights, dim=0)
adaln_linear_biases=torch.cat(adaln_linear_biases, dim=0)
in_features=adaln_linear_weights.shape[1]
out_features=adaln_linear_weights.shape[0]
model.adaln_linear_single=torch.nn.Linear(
in_features, out_features, bias=True, device=adaln_linear_weights.device, dtype=adaln_linear_weights.dtype
)
model.adaln_linear_single.weight.copy_(adaln_linear_weights)
model.adaln_linear_single.bias.copy_(adaln_linear_biases)
defprepare_clip_embeddings(
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
prompt: str,
device: torch.device,
dtype: torch.dtype,
max_length: int=77,
) ->torch.Tensor:
prompt= [prompt]
text_inputs=tokenizer(
prompt,
padding="max_length",
max_length=max_length,
truncation=True,
return_overflowing_tokens=False,
return_length=False,
return_tensors="pt",
)
text_input_ids=text_inputs.input_idsprompt_embeds=text_encoder(text_input_ids.to(device), output_hidden_states=False)
prompt_embeds=prompt_embeds.pooler_output.to(dtype)
returnprompt_embedsdefprepare_t5_embeddings(
text_encoder: T5EncoderModel,
tokenizer: T5TokenizerFast,
prompt: str,
device: torch.device,
dtype: torch.dtype,
max_length: int=512,
enable_prompt_length_bucketing: bool=False,
) ->torch.Tensor:
prompt= [prompt]
text_inputs=tokenizer(
prompt,
padding="max_length",
max_length=max_length,
truncation=True,
return_length=False,
return_overflowing_tokens=False,
return_tensors="pt",
)
text_input_ids=text_inputs.input_idsifenable_prompt_length_bucketing:
attention_mask=text_inputs.attention_masknum_text_tokens=attention_mask.sum(dim=1).max().item()
max_length=min(
SUPPORTED_BUCKET_LENGTHS, key=lambdax: abs(x-num_text_tokens) ifx>=num_text_tokenselsefloat("inf")
)
text_input_ids=text_input_ids[:, :max_length]
prompt_embeds=text_encoder(text_input_ids.to(device), output_hidden_states=False)[0]
prompt_embeds=prompt_embeds.to(dtype)
returnprompt_embeds, max_length@functools.lru_cache(maxsize=8)defprepare_latent_image_ids(height: int, width: int, device: torch.device, dtype: torch.dtype):
latent_image_ids=torch.zeros(height, width, 3)
latent_image_ids[..., 1] =latent_image_ids[..., 1] +torch.arange(height)[:, None]
latent_image_ids[..., 2] =latent_image_ids[..., 2] +torch.arange(width)[None, :]
latent_image_id_height, latent_image_id_width, latent_image_id_channels=latent_image_ids.shapelatent_image_ids=latent_image_ids.reshape(
latent_image_id_height*latent_image_id_width, latent_image_id_channels
).contiguous()
returnlatent_image_ids.to(device=device, dtype=dtype)
defprecompute_guidance_embeds(transformer: FluxTransformer2DModel, device: torch.device, dtype: torch.dtype):
embeds= {}
forguidance_scaleinSUPPORTED_GUIDANCE_SCALES:
guidance=torch.full([1], guidance_scale, device=device, dtype=torch.float32)
guidance=transformer.time_text_embed.guidance_embedder(
transformer.time_text_embed.time_proj(guidance*1000.0).to(dtype)
)
embeds[f"{guidance_scale:.1f}"] =guidancereturnembedsdefprecompute_timestep_embeds(transformer: FluxTransformer2DModel, device: torch.device, dtype: torch.dtype):
embeds= {}
image_seq_len=LATENT_HEIGHT*LATENT_WIDTHfornum_inference_stepsinrange(MIN_INFERENCE_STEPS, MAX_INFERENCE_STEPS+1):
sigmas=np.linspace(1.0, 1/num_inference_steps, num_inference_steps)
mu=B+image_seq_len*Msigmas=math.exp(mu) / (math.exp(mu) + (1/sigmas-1) **1.0)
sigmas=torch.from_numpy(sigmas).to(device, dtype=torch.float32)
sigmas=torch.cat([sigmas, sigmas.new_zeros(1)])
timesteps= (sigmas*1000.0).to(dtype)
temb=transformer.time_text_embed.time_proj(timesteps)
temb=transformer.time_text_embed.timestep_embedder(temb.to(dtype))
embeds[num_inference_steps] = (sigmas, temb)
returnembedsdefprecompute_embeds(transformer: FluxTransformer2DModel, device: torch.device, dtype: torch.dtype, save_dir: str):
save_dir=pathlib.Path(save_dir)
save_dir.mkdir(parents=True, exist_ok=True)
guidance_path=save_dir/"guidance_embeds.pt"timestep_path=save_dir/"timestep_embeds.pt"ifguidance_path.exists():
guidance_embeds=torch.load(guidance_path, map_location=device, weights_only=True)
print(f'Loaded precomputed guidance embeddings from "{guidance_path.as_posix()}"')
else:
guidance_embeds=precompute_guidance_embeds(transformer, device, dtype)
ifcp_options.meshisNoneorcp_options.mesh._flatten().get_local_rank() ==0:
torch.save(guidance_embeds, guidance_path.as_posix())
print(f'Precomputed guidance embeddings saved to "{save_dir.as_posix()}"')
iftimestep_path.exists():
timestep_embeds=torch.load(timestep_path, map_location=device, weights_only=True)
print(f'Loaded precomputed timestep embeddings from "{timestep_path.as_posix()}"')
else:
timestep_embeds=precompute_timestep_embeds(transformer, device, dtype)
ifcp_options.meshisNoneorcp_options.mesh._flatten().get_local_rank() ==0:
torch.save(timestep_embeds, timestep_path.as_posix())
print(f'Precomputed timestep embeddings saved to "{save_dir.as_posix()}"')
returnguidance_embeds, timestep_embeds@torch.compiledefpointwise_add3_silu(x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) ->torch.Tensor:
returntorch.nn.functional.silu(x+y+z)
defcapture_cudagraph(
model: FluxTransformer2DModel,
latents: torch.Tensor,
encoder_hidden_states: torch.Tensor,
conditioning: torch.Tensor,
image_rotary_emb: Tuple[torch.Tensor, torch.Tensor],
dt: torch.Tensor,
):
print("Warming up CUDAGraph capture")
s=torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
withtorch.cuda.stream(s):
for_inrange(2):
_=model(
hidden_states=latents,
encoder_hidden_states=encoder_hidden_states,
conditioning=conditioning,
image_rotary_emb=image_rotary_emb,
dt=dt,
)
torch.cuda.current_stream().wait_stream(s)
print("Capturing CUDAGraph")
static_latents=latents.clone()
static_conditioning=conditioning.clone()
static_dt=dt.clone()
graph=torch.cuda.CUDAGraph()
withtorch.cuda.graph(graph):
static_x=model(
hidden_states=static_latents,
encoder_hidden_states=encoder_hidden_states,
conditioning=static_conditioning,
image_rotary_emb=image_rotary_emb,
dt=static_dt,
)
returngraph, static_latents, static_conditioning, static_dt, static_x@torch.inference_mode()defmain(
model_id: str,
prompt: str,
height: int,
width: int,
num_inference_steps: int,
guidance_scale: float,
compile_mode: str,
output_file: str,
cache_dir: Optional[str],
enable_cudagraph: bool,
enable_prompt_length_bucketing: bool,
enable_profiling: bool,
working_dir: str,
seed: int,
):
device="cuda"dtype=torch.bfloat16# Load the model componentstransformer=FluxTransformer2DModel.from_pretrained(
model_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=dtype
)
text_encoder=CLIPTextModel.from_pretrained(
model_id, subfolder="text_encoder", cache_dir=cache_dir, torch_dtype=dtype
)
text_encoder_2=T5EncoderModel.from_pretrained(
model_id, subfolder="text_encoder_2", cache_dir=cache_dir, torch_dtype=dtype
)
tokenizer=CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer", cache_dir=cache_dir)
tokenizer_2=T5TokenizerFast.from_pretrained(model_id, subfolder="tokenizer_2", cache_dir=cache_dir)
vae=AutoencoderKL.from_pretrained(model_id, subfolder="vae", cache_dir=cache_dir, torch_dtype=dtype)
image_processor=VaeImageProcessor(
vae_scale_factor=SPATIAL_COMPRESSION_RATIO*PIXEL_UNSHUFFLING_DOWNSAMPLING_FACTOR
)
fuse_qkv_(transformer)
fuse_adaln_linear_(transformer)
[x.to(device) forxin (transformer, text_encoder, text_encoder_2, vae)]
vae.to(memory_format=torch.channels_last)
ifcompile_mode!="none":
transformer=torch.compile(transformer, mode=compile_mode, fullgraph=True, dynamic=True)
# text_encoder = torch.compile(text_encoder, mode="default", fullgraph=True, dynamic=True)# text_encoder_2 = torch.compile(text_encoder_2, mode="default", fullgraph=True, dynamic=True)# We don't compile the VAE due to the implementation calling into non-traceable code paths# vae.decode = torch.compile(vae.decode, mode="default", fullgraph=True, dynamic=True)# Latent, text, guidance and timestep conditioning preparationbatch_size=1patch_size=transformer.config.patch_sizelatent_height=height// (SPATIAL_COMPRESSION_RATIO*patch_size) //PIXEL_UNSHUFFLING_DOWNSAMPLING_FACTORlatent_width=width// (SPATIAL_COMPRESSION_RATIO*patch_size) //PIXEL_UNSHUFFLING_DOWNSAMPLING_FACTORgenerator=torch.Generator(device=device).manual_seed(seed)
guidance_embeds, timestep_embeds=precompute_embeds(transformer, device, dtype, working_dir)
latents=torch.randn(
(batch_size, latent_height*latent_width, transformer.config.in_channels),
dtype=dtype,
device=device,
generator=generator,
)
pooled_projections=prepare_clip_embeddings(text_encoder, tokenizer, prompt, device, dtype)
encoder_hidden_states, num_text_tokens=prepare_t5_embeddings(
text_encoder_2, tokenizer_2, prompt, device, dtype, T5_SEQUENCE_LENGTH, enable_prompt_length_bucketing
)
# <precompute>guidance_conditioning=guidance_embeds[f"{guidance_scale:.1f}"]
sigmas, timestep_conditioning=timestep_embeds[num_inference_steps]
pooled_projections=transformer.time_text_embed.text_embedder(pooled_projections)
encoder_hidden_states=transformer.context_embedder(encoder_hidden_states)
# </precompute>img_ids=prepare_latent_image_ids(latent_height, latent_width, device=device, dtype=dtype)
txt_ids=torch.zeros(num_text_tokens, 3).to(device=device, dtype=dtype)
ifcp_options.meshisnotNone:
# Note: clone seems to be a must here otherwise there is a recompilation related to storage offsets (which# tells you to use torch._dynamo.decorators.mark_unbacked) /shrugimg_ids=EquipartitionSharder.shard(img_ids, dim=0, mesh=cp_options._flattened_mesh).clone()
txt_ids=EquipartitionSharder.shard(txt_ids, dim=0, mesh=cp_options._flattened_mesh).clone()
latents=EquipartitionSharder.shard(latents, dim=1, mesh=cp_options._flattened_mesh).clone()
encoder_hidden_states=EquipartitionSharder.shard(encoder_hidden_states, dim=1, mesh=cp_options._flattened_mesh).clone()
ids=torch.cat([txt_ids, img_ids], dim=0).float()
image_rotary_emb=transformer.pos_embed(ids)
dt=sigmas[1:] -sigmas[:-1]
print("Warming up the model")
for_inrange(2):
conditioning=pointwise_add3_silu(timestep_conditioning[0, :], guidance_conditioning, pooled_projections)
_=transformer(
hidden_states=latents,
encoder_hidden_states=encoder_hidden_states,
conditioning=conditioning,
image_rotary_emb=image_rotary_emb,
dt=dt[0],
)
torch.cuda.synchronize()
start_event=torch.cuda.Event(enable_timing=True)
end_event=torch.cuda.Event(enable_timing=True)
context= (
profile(activities=[ProfilerActivity.CUDA, ProfilerActivity.CPU], with_flops=True)
ifenable_profilingelsecontextlib.nullcontext()
)
ifnotenable_cudagraph:
withcontextasctx:
start_event.record()
foriinrange(num_inference_steps):
conditioning=pointwise_add3_silu(timestep_conditioning[i, :], guidance_conditioning, pooled_projections)
latents=transformer(
hidden_states=latents,
encoder_hidden_states=encoder_hidden_states,
conditioning=conditioning,
image_rotary_emb=image_rotary_emb,
dt=dt[i],
)
end_event.record()
torch.cuda.synchronize()
else:
conditioning=pointwise_add3_silu(timestep_conditioning[0, :], guidance_conditioning, pooled_projections)
graph, static_latents, static_conditioning, static_dt, static_x=capture_cudagraph(
transformer,
latents,
encoder_hidden_states,
conditioning,
image_rotary_emb,
dt[0],
)
withcontextasctx:
start_event.record()
static_x.copy_(latents)
foriinrange(num_inference_steps):
conditioning=pointwise_add3_silu(timestep_conditioning[i, :], guidance_conditioning, pooled_projections)
torch._foreach_copy_(
(static_latents, static_conditioning, static_dt),
(static_x, conditioning, dt[i]),
non_blocking=True,
)
graph.replay()
end_event.record()
torch.cuda.synchronize()
latents=static_xtotal_time=start_event.elapsed_time(end_event) /1000.0ifcp_options.meshisnotNone:
latents=EquipartitionSharder.unshard(latents, dim=1, mesh=cp_options._flattened_mesh)
ifcp_options.meshisNoneorcp_options.mesh._flatten().get_local_rank() ==0:
print(f"time: {total_time:.2f}s")
ifenable_profiling:
print(ctx.key_averages().table(sort_by="self_cuda_time_total", row_limit=-1))
ctx.export_chrome_trace("dump_benchmark_flux.json")
latents=latents.reshape(batch_size, latent_height, latent_width, -1, 2, 2)
latents=latents.permute(0, 3, 1, 4, 2, 5)
latents=latents.flatten(4, 5).flatten(2, 3)
latents=latents.to(memory_format=torch.channels_last)
latents= (latents/vae.config.scaling_factor) +vae.config.shift_factorimage=vae.decode(latents, return_dict=False)[0]
image=image_processor.postprocess(image, output_type="pil")[0]
image.save(output_file)
classEquipartitionSharder:
@classmethoddefshard(cls, tensor: torch.Tensor, dim: int, mesh: torch.distributed.device_mesh.DeviceMesh) ->torch.Tensor:
asserttensor.size()[dim] %mesh.size() ==0# The following is not fullgraph compatible with Dynamo (fails in DeviceMesh.get_rank)# return tensor.chunk(mesh.size(), dim=dim)[mesh.get_rank()]returntensor.chunk(mesh.size(), dim=dim)[torch.distributed.get_rank(mesh.get_group())]
@classmethoddefunshard(cls, tensor: torch.Tensor, dim: int, mesh: torch.distributed.device_mesh.DeviceMesh) ->torch.Tensor:
tensor=tensor.contiguous()
tensor=funcol.all_gather_tensor(tensor, dim, group=mesh.get_group())
returntensor# Reference:# - https://github.com/pytorch/pytorch/blob/f58a680d09e13658a52c6ba05c63c15759846bcc/torch/distributed/_functional_collectives.py#L827# - https://github.com/pytorch/pytorch/blob/f58a680d09e13658a52c6ba05c63c15759846bcc/torch/distributed/_functional_collectives.py#L246# For fullgraph=True tracing compatibility (since FakeTensor does not have a `wait` method):def_wait_tensor(tensor):
ifisinstance(tensor, funcol.AsyncCollectiveTensor):
tensor=tensor.wait()
returntensordef_all_to_all_single(x: torch.Tensor, group) ->torch.Tensor:
shape=x.shape# HACK: We need to flatten because despite making tensors contiguous, torch single-file-ization# to benchmark triton codegen fails somewhere:# buf25 = torch.ops._c10d_functional.all_to_all_single.default(buf24, [1, 1], [1, 1], '3')# ValueError: Tensors must be contiguousx=x.flatten()
x=funcol.all_to_all_single(x, None, None, group)
x=x.reshape(shape)
x=_wait_tensor(x)
returnxdef_templated_ring_attention(query, key, value):
ring_mesh=cp_options.mesh["ring"]
rank=cp_options._ring_local_rankworld_size=cp_options.ring_degreeifworld_size==1:
returncp_options.attention_op(query, key, value)
next_rank= (rank+1) %world_sizeprev_out=prev_lse=Nonekv_buffer=torch.cat([key.flatten(), value.flatten()]).contiguous()
kv_buffer=funcol.all_gather_tensor(kv_buffer, gather_dim=0, group=ring_mesh.get_group())
kv_buffer=kv_buffer.chunk(world_size)
foriinrange(world_size):
ifi>0:
kv=kv_buffer[next_rank]
key=kv[:key.numel()].reshape_as(key)
value=kv[key.numel():].reshape_as(value)
next_rank= (next_rank+1) %world_sizeout, lse=cp_options.attention_op(query, key, value)
ifcp_options.convert_to_fp32:
out=out.to(torch.float32)
lse=lse.to(torch.float32)
lse=lse.unsqueeze(-1)
ifprev_outisnotNone:
out=prev_out-torch.nn.functional.sigmoid(lse-prev_lse) * (prev_out-out)
lse=prev_lse-torch.nn.functional.logsigmoid(prev_lse-lse)
prev_out=outprev_lse=lseout=out.to(query.dtype)
lse=lse.squeeze(-1)
returnout, lsedef_templated_ulysses_attention(query, key, value, *, return_lse: bool=False):
ulysses_mesh=cp_options.mesh["ulysses"]
world_size=cp_options.ulysses_degreegroup=ulysses_mesh.get_group()
ifworld_size==1:
returncp_options.attention_op(query, key, value)
B, S_LOCAL, H, D=query.shapeH_LOCAL=H//world_sizequery, key, value= (
x.reshape(B, S_LOCAL, world_size, H_LOCAL, D).permute(2, 1, 0, 3, 4).clone()
forxin (query, key, value)
)
query, key, value= (
_all_to_all_single(x, group)
forxin (query, key, value)
)
query, key, value= (
x.flatten(0, 1).permute(1, 0, 2, 3).contiguous()
forxin (query, key, value)
)
out, lse=cp_options.attention_op(query, key, value)
out=out.reshape(B, world_size, S_LOCAL, H_LOCAL, D).permute(1, 3, 0, 2, 4).contiguous()
out=_all_to_all_single(out, group)
out=out.flatten(0, 1).permute(1, 2, 0, 3).contiguous()
ifreturn_lse:
lse=lse.reshape(B, world_size, S_LOCAL, H_LOCAL).permute(1, 3, 0, 2).contiguous()
lse=_all_to_all_single(lse, group)
lse=lse.flatten(0, 1).permute(1, 2, 0).contiguous()
else:
lse=Nonereturnout, lse# TODO: currently produces incorrect results (for example, with CP=4, ring=2, ulysses=2, half output is our expected image,# and other half is some completely different)# def _templated_unified_attention(query, key, value):# ulysses_mesh = cp_options.mesh["ulysses"]# ulysses_size = ulysses_mesh.size()# ulysses_group = ulysses_mesh.get_group()# B, S_LOCAL, H, D = query.shape# H_LOCAL = H // ulysses_size# query, key, value = (# x.reshape(B, S_LOCAL, ulysses_size, H_LOCAL, D).permute(2, 1, 0, 3, 4).contiguous()# for x in (query, key, value)# )# query, key, value = (# wait_tensor(funcol.all_to_all_single(x, None, None, group=ulysses_group))# for x in (query, key, value)# )# query, key, value = (# x.flatten(0, 1).permute(1, 0, 2, 3).contiguous()# for x in (query, key, value)# )# out, lse = _templated_ring_attention(query, key, value)# out = out.reshape(B, ulysses_size, S_LOCAL, H_LOCAL, D).permute(1, 3, 0, 2, 4).contiguous()# lse = lse.reshape(B, ulysses_size, S_LOCAL, H_LOCAL).permute(1, 3, 0, 2).contiguous()# out = wait_tensor(funcol.all_to_all_single(out, None, None, group=ulysses_group))# lse = wait_tensor(funcol.all_to_all_single(lse, None, None, group=ulysses_group))# out = out.flatten(0, 1).permute(1, 2, 0, 3).contiguous()# lse = lse.flatten(0, 1).permute(1, 2, 0).contiguous()# return out, lse# For fullgraph=True tracing to be compatible@torch.library.custom_op("flash_attn_3::_flash_attn_forward_original", mutates_args=(), device_types="cuda")def_wrapped_flash_attn_3_original(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) ->Tuple[torch.Tensor, torch.Tensor]:
out, lse=flash_attn_3_func(query, key, value)
lse=lse.permute(0, 2, 1)
returnout, lse@torch.library.register_fake("flash_attn_3::_flash_attn_forward_original")def_(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) ->Tuple[torch.Tensor, torch.Tensor]:
batch_size, seq_len, num_heads, head_dim=query.shapelse_shape= (batch_size, seq_len, num_heads)
returntorch.empty_like(query), query.new_empty(lse_shape)
@torch.library.custom_op("flash_attn_3::_flash_attn_forward_hf", mutates_args=(), device_types="cuda")def_wrapped_flash_attn_3_hf(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) ->torch.Tensor:
out, lse=flash_attn_3_hf.flash_attn_func(query, key, value, causal=False)
returnout@torch.library.register_fake("flash_attn_3::_flash_attn_forward_hf")def_(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) ->torch.Tensor:
returntorch.empty_like(query)
def_attention_torch_cudnn(query, key, value):
query, key, value= (x.transpose(1, 2).contiguous() forxin (query, key, value))
out, lse, cum_seq_q, cum_seq_k, max_q, max_k, philox_seed, philox_offset, debug_attn_mask= (
torch.ops.aten._scaled_dot_product_cudnn_attention(
query=query,
key=key,
value=value,
attn_bias=None,
compute_log_sumexp=True,
)
)
out=out.transpose(1, 2).contiguous()
lse=lse.transpose(1, 2).contiguous()
returnout, lsedef_attention_flash_attn_2(query, key, value):
out, lse, _=flash_attn_func(query, key, value, return_attn_probs=True)
lse=lse.permute(0, 2, 1)
returnout, lsedef_attention_flash_attn_3_original(query, key, value):
out=_wrapped_flash_attn_3_original(query, key, value)
returnoutdef_attention_flash_attn_3_hf(query, key, value):
out=_wrapped_flash_attn_3_hf(query, key, value)
returnoutdef_download_hf_flash_attn_3():
globalflash_attn_3_hfflash_attn_3_hf=get_kernel("kernels-community/flash-attn3")
defget_args():
DEFAULT_MODEL_ID="black-forest-labs/FLUX.1-dev"DEFAULT_PROMPT="The King of Hearts card transforms into a 3D hologram that appears to be made of cosmic energy. As the King emerges, stars and galaxies swirl around him, creating a sense of traveling through the universe. The King's attire is adorned with celestial patterns, and his crown is a glowing star cluster. The hologram floats in front of you, with the background shifting through different cosmic scenes, from nebulae to black holes. Atmosphere: Perfect for space-themed events, science fiction conventions, or futuristic tech expos."parser=argparse.ArgumentParser()
parser.add_argument("--model_id", type=str, default=DEFAULT_MODEL_ID)
parser.add_argument("--prompt", type=str, default=DEFAULT_PROMPT)
parser.add_argument("--height", type=int, default=1024)
parser.add_argument("--width", type=int, default=1024)
parser.add_argument("--num_inference_steps", type=int, default=28)
parser.add_argument("--guidance_scale", type=float, default=4.0)
parser.add_argument(
"--compile_mode",
type=str,
default="none",
choices=["none", "default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"],
)
parser.add_argument("--output_file", type=str, default="output.png")
parser.add_argument("--cache_dir", type=str, default=None)
parser.add_argument("--enable_fp32_rope", action="store_true")
parser.add_argument("--enable_cudagraph", action="store_true")
parser.add_argument("--enable_prompt_length_bucketing", action="store_true")
parser.add_argument("--attention_provider", type=str, default="cudnn", choices=["cudnn", "fa2", "fa3", "fa3_original"])
parser.add_argument("--ring_degree", type=int, default=1)
parser.add_argument("--ulysses_degree", type=int, default=1)
parser.add_argument("--enable_profiling", action="store_true")
parser.add_argument("--disable_tf32", action="store_true")
parser.add_argument("--disable_flags", action="store_true")
parser.add_argument("--working_dir", type=str, default="/tmp/flux_precomputation")
parser.add_argument("--seed", type=int, default=31337)
args=parser.parse_args()
returnargsdefsetup_config(args):
torch.manual_seed(args.seed)
ifargs.enable_fp32_rope:
globalROPE_PRECISIONROPE_PRECISION=torch.float32ifargs.enable_cudagraphandargs.compile_modenotin ["none", "default", "max-autotune-no-cudagraphs"]:
raiseValueError(
"Only compiled modes 'none', 'default', and 'max-autotune-no-cudagraphs' are supported with CUDAGraphs."
)
globalATTENTION_OPifargs.attention_provider=="cudnn":
ATTENTION_OP=_attention_torch_cudnnelifargs.attention_provider=="fa2":
ATTENTION_OP=_attention_flash_attn_2elifargs.attention_provider=="fa3":
_download_hf_flash_attn_3()
ATTENTION_OP=_attention_flash_attn_3_hfelifargs.attention_provider=="fa3_original":
ATTENTION_OP=_attention_flash_attn_3_originalelse:
assertFalseifargs.enable_profilingandargs.enable_cudagraph:
torch.profiler._utils._init_for_cuda_graphs()
ifnotargs.disable_tf32:
torch.backends.cuda.matmul.allow_tf32=Truetorch.backends.cudnn.allow_tf32=Trueifnotargs.disable_flags:
apply_flags()
globalMIN_INFERENCE_STEPS, MAX_INFERENCE_STEPSifargs.num_inference_steps<MIN_INFERENCE_STEPSorargs.num_inference_steps>MAX_INFERENCE_STEPS:
raiseValueError(f"`num_inference_steps` must be equal or between {MIN_INFERENCE_STEPS} and {MAX_INFERENCE_STEPS}.")
defsetup_distributed(ring_degree: int, ulysses_degree: int, compile_mode: str):
globalATTENTION_OPdist.init_process_group("nccl")
torch.cuda.set_device(torch.device("cuda", dist.get_rank()))
ifring_degree*ulysses_degree!=dist.get_world_size():
raiseValueError(f"`ring_degree * ulysses_degree` must equal the world size {dist.get_world_size()}.")
mesh_names= ["ring", "ulysses"]
mesh_dims= [ring_degree, ulysses_degree]
mesh=dist.device_mesh.init_device_mesh("cuda", mesh_dims, mesh_dim_names=mesh_names)
cp_options.ring_degree=ring_degreecp_options.ulysses_degree=ulysses_degreecp_options.mesh=meshcp_options.convert_to_fp32=Truecp_options.attention_op=ATTENTION_OPcp_options._flattened_mesh=mesh._flatten()
cp_options._ring_mesh=mesh["ring"]
cp_options._ulysses_mesh=mesh["ulysses"]
cp_options._ring_local_rank=cp_options._ring_mesh.get_local_rank()
cp_options._ulysses_local_rank=cp_options._ulysses_mesh.get_local_rank()
ifring_degree>1andulysses_degree>1:
raiseValueError("The current implementation is incorrect for unified attention and needs to be fixed.")
# cp_options.mode = "unified"# ATTENTION_OP = _templated_unified_attentionelifulysses_degree>1:
cp_options.mode="ulysses"ATTENTION_OP=_templated_ulysses_attentionelse:
cp_options.mode="ring"ATTENTION_OP=_templated_ring_attentionifcompile_mode!="none":
torch._dynamo.config.suppress_errors=Truetorch._inductor.config.reorder_for_compute_comm_overlap=Trueif__name__=="__main__":
args=get_args()
try:
setup_config(args)
setup_distributed(args.ring_degree, args.ulysses_degree, args.compile_mode)
main(
model_id=args.model_id,
prompt=args.prompt,
height=args.height,
width=args.width,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
compile_mode=args.compile_mode,
output_file=args.output_file,
cache_dir=args.cache_dir,
enable_cudagraph=args.enable_cudagraph,
enable_prompt_length_bucketing=args.enable_prompt_length_bucketing,
enable_profiling=args.enable_profiling,
working_dir=args.working_dir,
seed=args.seed,
)
exceptExceptionase:
print(f"An error occurred: {e}")
ifdist.is_initialized():
torch.distributed.breakpoint()
raisefinally:
ifdist.is_initialized():
dist.destroy_process_group()

A100	H100

TODO: link to blog post

Explanation

Each model should define a _cp_plan attribute that contains information on how to shard/gather tensors at different stages of the forward. Let's try to understand with an example using QwenImage:

_cp_plan= {
"": {
"hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
"encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
"encoder_hidden_states_mask": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False),
},
"pos_embed": {
0: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True),
1: ContextParallelInput(split_dim=0, expected_dims=2, split_output=True),
},
"proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
}

The dictionary is a set of module names mapped to their respective CP plan. The inputs/outputs of layers will be split/gathered according to this at the respective module level. Here, the following happens:

• "": we specify that we want to split the various inputs across the sequence dim in the pre-forward hook (i.e. before the actual forward logic of the QwenImageTransformer2DModel is run, we will split the inputs)
• "pos_embed": we specify that we want to split the outputs of the RoPE layer. Since there are two outputs (imag & text freqs), we can individually specify how they should be split
• "proj_out": before returning to the user, we gather the entire sequence on each rank in the post-forward hook (after the linear layer forward has run).

ContextParallelInput: specifies how to split the input tensor in the pre-forward or post-forward hook of the layer it is attached to

ContextParallelOutput: specifies how to gather the input tensor in the post-forward hook in the layer it is attached to

[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.

[DN6]

Great work! Just a few questions and nits.

[DN6]

Hmm wondering if we need this additional context manager. I think the name implies that it's parallelizing the components, when it's really a validation step.

[a-r-r-o-w]

Actually, this function is what sets the dispatcher to perform context parallel templated attention instead of following the normal attention path in non-CP case

diffusers/src/diffusers/models/_modeling_parallel.py

Line 171 in bb07c04

_AttentionBackendRegistry._parallel_config=model_or_pipeline._internal_parallel_config

Without this, we have to hook into perform the assignment in the pre-forward hook of the model. But that has compatibility issues with torch dynamo (tracing feels with a setattr/getattr-related error. It was our previous approach in this commit, but it had many compatibility issues for distributed training. Instead, explicitly doing this outside the forward is ideal for setting up all the required information about parallelism.

[DN6]

Hmm but setting parallel_config in this way leads to this issue no?

[DN6]

Great work as always @a-r-r-o-w. Good to merge once tests are passing.

[a-r-r-o-w]

Thanks @DN6! I'm AFK from my personal laptop so can't make changes for another 2-3 days. Sorry about the delay!

[a-r-r-o-w]

@DN6 I'm unable to push any changes to this branch since it's on the official repo instead of the my personal fork. I think if you added an entry to _toctree.yaml, it should fix the failing test

[sayakpaul]

I will follow suit and make it work with FA3 Hub as well.

[sayakpaul]

Just pushed an entry to the TOC. Cc: @stevhliu for now, I have added it to "Inference optimization" section. But feel free to change that in your documentation PR.

@a-r-r-o-w sorry for the delay here. We should be able to merge this now. Thanks for setting the foundations here :)