diff --git a/src/maxdiffusion/configs/base_flux2klein.yml b/src/maxdiffusion/configs/base_flux2klein.yml new file mode 100644 index 000000000..7c39ca2da --- /dev/null +++ b/src/maxdiffusion/configs/base_flux2klein.yml @@ -0,0 +1,278 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# This sentinel is a reminder to choose a real run name. +run_name: 'flux2klein_test_run' + +metrics_file: "" # for testing, local file that stores scalar metrics. If empty, no metrics are written. +# If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/ +write_metrics: True + +timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written. +write_timing_metrics: True + +gcs_metrics: False +# If true save config to GCS in {base_output_directory}/{run_name}/ +save_config_to_gcs: False +log_period: 100 + +pretrained_model_name_or_path: 'black-forest-labs/FLUX.2-klein-4B' +clip_model_name_or_path: 'ariG23498/clip-vit-large-patch14-text-flax' +t5xxl_model_name_or_path: 'ariG23498/t5-v1-1-xxl-flax' + +# Flux params +flux_name: "flux2klein" +scale_shift_order: "scale_shift" +use_latents: False +max_sequence_length: 512 +time_shift: True +base_shift: 0.5 +max_shift: 1.15 +# offloads t5 encoder after text encoding to save memory. +offload_encoders: True + + +unet_checkpoint: '' +revision: 'refs/pr/95' +# This will convert the weights to this dtype. +# When running inference on TPUv5e, use weights_dtype: 'bfloat16' +weights_dtype: 'bfloat16' +# This sets the layer's dtype in the model. Ex: nn.Dense(dtype=activations_dtype) +activations_dtype: 'bfloat16' + +# matmul and conv precision from https://jax.readthedocs.io/en/latest/jax.lax.html#jax.lax.Precision +# Options are "DEFAULT", "HIGH", "HIGHEST" +# fp32 activations and fp32 weights with HIGHEST will provide the best precision +# at the cost of time. +precision: "DEFAULT" + +# if False state is not jitted and instead replicate is called. This is good for debugging on single host +# It must be True for multi-host. +jit_initializers: True + +# Set true to load weights from pytorch +from_pt: True +split_head_dim: True +attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te +# If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens. +# Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster. +# However, when padding tokens are significant, this will lead to worse quality and should be set to True. +mask_padding_tokens: True +# Maxdiffusion has 2 types of attention sharding strategies: +# 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention) +# 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded +# in cross attention q. +attention_sharding_uniform: True + +flash_block_sizes: {} +# GroupNorm groups +norm_num_groups: 32 + +# If train_new_flux, flux weights will be randomly initialized to train flux from scratch +# else they will be loaded from pretrained_model_name_or_path +train_new_flux: False + +# train text_encoder - Currently not supported for SDXL +train_text_encoder: False +text_encoder_learning_rate: 4.25e-6 + +# https://arxiv.org/pdf/2305.08891.pdf +snr_gamma: -1.0 + +timestep_bias: { + # a value of later will increase the frequence of the model's final training steps. + # none, earlier, later, range + strategy: "none", + # multiplier for bias, a value of 2.0 will double the weight of the bias, 0.5 will halve it. + multiplier: 1.0, + # when using strategy=range, the beginning (inclusive) timestep to bias. + begin: 0, + # when using strategy=range, the final step (inclusive) to bias. + end: 1000, + # portion of timesteps to bias. + # 0.5 will bias one half of the timesteps. Value of strategy determines + # whether the biased portions are in the earlier or later timesteps. + portion: 0.25 +} + +# Override parameters from checkpoints's scheduler. +diffusion_scheduler_config: { + _class_name: 'FlaxEulerDiscreteScheduler', + prediction_type: 'epsilon', + rescale_zero_terminal_snr: False, + timestep_spacing: 'trailing' +} + +# Output directory +# Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/" +base_output_directory: "" + +# Hardware +hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu' +skip_jax_distributed_system: False + +# Parallelism +mesh_axes: ['data', 'fsdp', 'context', 'tensor'] + +# batch : batch dimension of data and activations +# hidden : +# embed : attention qkv dense layer hidden dim named as embed +# heads : attention head dim = num_heads * head_dim +# length : attention sequence length +# temb_in : dense.shape[0] of resnet dense before conv +# out_c : dense.shape[1] of resnet dense before conv +# out_channels : conv.shape[-1] activation +# keep_1 : conv.shape[0] weight +# keep_2 : conv.shape[1] weight +# conv_in : conv.shape[2] weight +# conv_out : conv.shape[-1] weight +logical_axis_rules: [ + ['batch', 'data'], + ['activation_batch', ['data','fsdp']], + ['activation_heads', 'tensor'], + ['activation_kv', 'tensor'], + ['mlp','tensor'], + ['embed','fsdp'], + ['heads', 'tensor'], + ['conv_batch', ['data','fsdp']], + ['out_channels', 'tensor'], + ['conv_out', 'fsdp'], + ] +data_sharding: [['data', 'fsdp', 'context', 'tensor']] + +# One axis for each parallelism type may hold a placeholder (-1) +# value to auto-shard based on available slices and devices. +# By default, product of the DCN axes should equal number of slices +# and product of the ICI axes should equal number of devices per slice. +dcn_data_parallelism: 1 # recommended DCN axis to be auto-sharded +dcn_fsdp_parallelism: -1 +dcn_context_parallelism: 1 +dcn_tensor_parallelism: 1 +ici_data_parallelism: 1 +ici_fsdp_parallelism: -1 # recommended ICI axis to be auto-sharded +ici_context_parallelism: 1 +ici_tensor_parallelism: 1 + +allow_split_physical_axes: False + +# Dataset +# Replace with dataset path or train_data_dir. One has to be set. +dataset_name: 'diffusers/pokemon-gpt4-captions' +train_split: 'train' +dataset_type: 'tfrecord' # Options: 'tfrecord', 'hf', 'tf', 'grain', 'synthetic' +cache_latents_text_encoder_outputs: True +dataset_save_location: '/tmp/pokemon-gpt4-captions_xl' +train_data_dir: '' +dataset_config_name: '' +jax_cache_dir: '/tmp/jax_cache' +hf_data_dir: '' +hf_train_files: '' +hf_access_token: '' +image_column: 'image' +caption_column: 'text' +resolution: 512 +center_crop: False +random_flip: False +tokenize_captions_num_proc: 4 +transform_images_num_proc: 4 +reuse_example_batch: False +enable_data_shuffling: True + +# checkpoint every number of samples, -1 means don't checkpoint. +checkpoint_every: -1 +# enables one replica to read the ckpt then broadcast to the rest +enable_single_replica_ckpt_restoring: False + +# Training loop +learning_rate: 1.e-5 +scale_lr: False +max_train_samples: -1 +# max_train_steps takes priority over num_train_epochs. +max_train_steps: 1500 +num_train_epochs: 1 +seed: 0 +output_dir: 'output/' +per_device_batch_size: 1 + +warmup_steps_fraction: 0.1 +learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps. + +# AdamW optimizer parameters +adam_b1: 0.9 # Exponential decay rate to track the first moment of past gradients. +adam_b2: 0.999 # Exponential decay rate to track the second moment of past gradients. +adam_eps: 1.e-8 # A small constant applied to denominator outside of the square root. +adam_weight_decay: 0 # AdamW Weight decay +opt_enable_grad_clipping: False +max_grad_value: 1.0 +opt_enable_grad_global_norm_clipping: False +max_grad_norm: 1.0 + +enable_profiler: False +skip_first_n_steps_for_profiler: 5 +profiler_steps: 10 +profiler: "" + +# Generation parameters +prompt: "A detailed vector illustration of a robotic hummingbird || A cinematic shot of a neon-lit cyberpunk street" +prompt_2: "A detailed vector illustration of a robotic hummingbird || A cinematic shot of a neon-lit cyberpunk street" +negative_prompt: "" +do_classifier_free_guidance: True +guidance_scale: 4.0 +guidance_rescale: 0.0 +num_inference_steps: 4 +save_final_checkpoint: False + +# SDXL Lightning parameters +lightning_from_pt: True +lightning_repo: "" +lightning_ckpt: "" + +# LoRA parameters +lora_config: { + lora_model_name_or_path: [], + weight_name: [], + adapter_name: [], + scale: [], + from_pt: [] +} + +enable_mllog: False + +#controlnet +controlnet_model_name_or_path: 'diffusers/controlnet-canny-sdxl-1.0' +controlnet_from_pt: True +controlnet_conditioning_scale: 0.5 +controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Google_%22G%22_logo.svg/1024px-Google_%22G%22_logo.svg.png' +quantization: '' +quantization_local_shard_count: -1 +use_qwix_quantization: False +compile_topology_num_slices: -1 # Number of target slices, set to a positive integer. + +# ML Diagnostics settings +enable_ml_diagnostics: False +profiler_gcs_path: "" +enable_ondemand_xprof: False + +# Specific additions for generate_flux2klein execution +height: 1024 +width: 1024 +batch_size: 4 +interactive: False + +# 4B Architecture Dimensions +depth: 20 # num_single_layers +num_double_layers: 5 +hidden_size: 3072 +num_attention_heads: 24 + diff --git a/src/maxdiffusion/configs/base_flux2klein_9B.yml b/src/maxdiffusion/configs/base_flux2klein_9B.yml new file mode 100644 index 000000000..e6cba5951 --- /dev/null +++ b/src/maxdiffusion/configs/base_flux2klein_9B.yml @@ -0,0 +1,278 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# This sentinel is a reminder to choose a real run name. +run_name: 'flux2klein_9b_test_run' + +metrics_file: "" # for testing, local file that stores scalar metrics. If empty, no metrics are written. +# If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/ +write_metrics: True + +timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written. +write_timing_metrics: True + +gcs_metrics: False +# If true save config to GCS in {base_output_directory}/{run_name}/ +save_config_to_gcs: False +log_period: 100 + +pretrained_model_name_or_path: 'black-forest-labs/FLUX.2-klein-9B' +clip_model_name_or_path: 'ariG23498/clip-vit-large-patch14-text-flax' +t5xxl_model_name_or_path: 'ariG23498/t5-v1-1-xxl-flax' + +# Flux params +flux_name: "flux2klein_9B" +scale_shift_order: "scale_shift" +use_latents: False +max_sequence_length: 512 +time_shift: True +base_shift: 0.5 +max_shift: 1.15 +# offloads t5 encoder after text encoding to save memory. +offload_encoders: True + + +unet_checkpoint: '' +revision: 'refs/pr/95' +# This will convert the weights to this dtype. +# When running inference on TPUv5e, use weights_dtype: 'bfloat16' +weights_dtype: 'bfloat16' +# This sets the layer's dtype in the model. Ex: nn.Dense(dtype=activations_dtype) +activations_dtype: 'bfloat16' + +# matmul and conv precision from https://jax.readthedocs.io/en/latest/jax.lax.html#jax.lax.Precision +# Options are "DEFAULT", "HIGH", "HIGHEST" +# fp32 activations and fp32 weights with HIGHEST will provide the best precision +# at the cost of time. +precision: "DEFAULT" + +# if False state is not jitted and instead replicate is called. This is good for debugging on single host +# It must be True for multi-host. +jit_initializers: True + +# Set true to load weights from pytorch +from_pt: True +split_head_dim: True +attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te +# If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens. +# Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster. +# However, when padding tokens are significant, this will lead to worse quality and should be set to True. +mask_padding_tokens: True +# Maxdiffusion has 2 types of attention sharding strategies: +# 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention) +# 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded +# in cross attention q. +attention_sharding_uniform: True + +flash_block_sizes: {} +# GroupNorm groups +norm_num_groups: 32 + +# If train_new_flux, flux weights will be randomly initialized to train flux from scratch +# else they will be loaded from pretrained_model_name_or_path +train_new_flux: False + +# train text_encoder - Currently not supported for SDXL +train_text_encoder: False +text_encoder_learning_rate: 4.25e-6 + +# https://arxiv.org/pdf/2305.08891.pdf +snr_gamma: -1.0 + +timestep_bias: { + # a value of later will increase the frequence of the model's final training steps. + # none, earlier, later, range + strategy: "none", + # multiplier for bias, a value of 2.0 will double the weight of the bias, 0.5 will halve it. + multiplier: 1.0, + # when using strategy=range, the beginning (inclusive) timestep to bias. + begin: 0, + # when using strategy=range, the final step (inclusive) to bias. + end: 1000, + # portion of timesteps to bias. + # 0.5 will bias one half of the timesteps. Value of strategy determines + # whether the biased portions are in the earlier or later timesteps. + portion: 0.25 +} + +# Override parameters from checkpoints's scheduler. +diffusion_scheduler_config: { + _class_name: 'FlaxEulerDiscreteScheduler', + prediction_type: 'epsilon', + rescale_zero_terminal_snr: False, + timestep_spacing: 'trailing' +} + +# Output directory +# Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/" +base_output_directory: "" + +# Hardware +hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu' +skip_jax_distributed_system: False + +# Parallelism +mesh_axes: ['data', 'fsdp', 'context', 'tensor'] + +# batch : batch dimension of data and activations +# hidden : +# embed : attention qkv dense layer hidden dim named as embed +# heads : attention head dim = num_heads * head_dim +# length : attention sequence length +# temb_in : dense.shape[0] of resnet dense before conv +# out_c : dense.shape[1] of resnet dense before conv +# out_channels : conv.shape[-1] activation +# keep_1 : conv.shape[0] weight +# keep_2 : conv.shape[1] weight +# conv_in : conv.shape[2] weight +# conv_out : conv.shape[-1] weight +logical_axis_rules: [ + ['batch', 'data'], + ['activation_batch', ['data','fsdp']], + ['activation_heads', 'tensor'], + ['activation_kv', 'tensor'], + ['mlp','tensor'], + ['embed','fsdp'], + ['heads', 'tensor'], + ['conv_batch', ['data','fsdp']], + ['out_channels', 'tensor'], + ['conv_out', 'fsdp'], + ] +data_sharding: [['data', 'fsdp', 'context', 'tensor']] + +# One axis for each parallelism type may hold a placeholder (-1) +# value to auto-shard based on available slices and devices. +# By default, product of the DCN axes should equal number of slices +# and product of the ICI axes should equal number of devices per slice. +dcn_data_parallelism: 1 # recommended DCN axis to be auto-sharded +dcn_fsdp_parallelism: -1 +dcn_context_parallelism: 1 +dcn_tensor_parallelism: 1 +ici_data_parallelism: 1 +ici_fsdp_parallelism: -1 # recommended ICI axis to be auto-sharded +ici_context_parallelism: 1 +ici_tensor_parallelism: 1 + +allow_split_physical_axes: False + +# Dataset +# Replace with dataset path or train_data_dir. One has to be set. +dataset_name: 'diffusers/pokemon-gpt4-captions' +train_split: 'train' +dataset_type: 'tfrecord' # Options: 'tfrecord', 'hf', 'tf', 'grain', 'synthetic' +cache_latents_text_encoder_outputs: True +dataset_save_location: '/tmp/pokemon-gpt4-captions_xl' +train_data_dir: '' +dataset_config_name: '' +jax_cache_dir: '/tmp/jax_cache' +hf_data_dir: '' +hf_train_files: '' +hf_access_token: '' +image_column: 'image' +caption_column: 'text' +resolution: 512 +center_crop: False +random_flip: False +tokenize_captions_num_proc: 4 +transform_images_num_proc: 4 +reuse_example_batch: False +enable_data_shuffling: True + +# checkpoint every number of samples, -1 means don't checkpoint. +checkpoint_every: -1 +# enables one replica to read the ckpt then broadcast to the rest +enable_single_replica_ckpt_restoring: False + +# Training loop +learning_rate: 1.e-5 +scale_lr: False +max_train_samples: -1 +# max_train_steps takes priority over num_train_epochs. +max_train_steps: 1500 +num_train_epochs: 1 +seed: 0 +output_dir: 'output/' +per_device_batch_size: 1 + +warmup_steps_fraction: 0.1 +learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps. + +# AdamW optimizer parameters +adam_b1: 0.9 # Exponential decay rate to track the first moment of past gradients. +adam_b2: 0.999 # Exponential decay rate to track the second moment of past gradients. +adam_eps: 1.e-8 # A small constant applied to denominator outside of the square root. +adam_weight_decay: 0 # AdamW Weight decay +opt_enable_grad_clipping: False +max_grad_value: 1.0 +opt_enable_grad_global_norm_clipping: False +max_grad_norm: 1.0 + +enable_profiler: False +skip_first_n_steps_for_profiler: 5 +profiler_steps: 10 +profiler: "" + +# Generation parameters +prompt: "A detailed vector illustration of a robotic hummingbird || A cinematic shot of a neon-lit cyberpunk street" +prompt_2: "A detailed vector illustration of a robotic hummingbird || A cinematic shot of a neon-lit cyberpunk street" +negative_prompt: "" +do_classifier_free_guidance: True +guidance_scale: 4.0 +guidance_rescale: 0.0 +num_inference_steps: 4 +save_final_checkpoint: False + +# SDXL Lightning parameters +lightning_from_pt: True +lightning_repo: "" +lightning_ckpt: "" + +# LoRA parameters +lora_config: { + lora_model_name_or_path: [], + weight_name: [], + adapter_name: [], + scale: [], + from_pt: [] +} + +enable_mllog: False + +#controlnet +controlnet_model_name_or_path: 'diffusers/controlnet-canny-sdxl-1.0' +controlnet_from_pt: True +controlnet_conditioning_scale: 0.5 +controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Google_%22G%22_logo.svg/1024px-Google_%22G%22_logo.svg.png' +quantization: '' +quantization_local_shard_count: -1 +use_qwix_quantization: False +compile_topology_num_slices: -1 # Number of target slices, set to a positive integer. + +# ML Diagnostics settings +enable_ml_diagnostics: False +profiler_gcs_path: "" +enable_ondemand_xprof: False + +# Specific additions for generate_flux2klein execution +height: 1024 +width: 1024 +batch_size: 4 +interactive: False + +# 9B Architecture Dimensions +depth: 24 # num_single_layers +num_double_layers: 8 +hidden_size: 4096 +num_attention_heads: 32 + diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 0ddbed625..edc9f4f7b 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -15,7 +15,6 @@ import contextlib import functools import math -import os from typing import Optional, Callable, Tuple, Any, Dict import flax.linen as nn from flax import nnx @@ -348,42 +347,6 @@ def convert_to_tokamax_splash_config( ) -def _extract_custom_block_sizes(flash_block_sizes): - """Pulls custom-kernel block sizes out of the (dict or BlockSizes-like) config. - - Mirrors the extraction used by the `ulysses_custom` path so the custom ring - kernel honors the same `flash_block_sizes={...}` knobs. - """ - bq = 4864 - bkv = 1024 - bkv_compute = 1024 - bkv_compute_in = 1024 - heads_per_tile = 1 - vmem_limit_bytes = None - if flash_block_sizes is not None: - if isinstance(flash_block_sizes, dict): - get = flash_block_sizes.get - bq = get("block_q", bq) - bkv = get("block_kv", bkv) - bkv_compute = get("block_kv_compute", bkv_compute) - bkv_compute_in = get("block_kv_compute_in", bkv_compute_in) - heads_per_tile = get("heads_per_tile", heads_per_tile) - vmem_limit_bytes = get("vmem_limit_bytes", vmem_limit_bytes) - else: - bq = getattr(flash_block_sizes, "block_q", bq) - bkv = getattr(flash_block_sizes, "block_kv", bkv) - bkv_compute = getattr(flash_block_sizes, "block_kv_compute", bkv_compute) - bkv_compute_in = getattr(flash_block_sizes, "block_kv_compute_in", bkv_compute_in) - heads_per_tile = getattr(flash_block_sizes, "heads_per_tile", heads_per_tile) - vmem_limit_bytes = getattr(flash_block_sizes, "vmem_limit_bytes", vmem_limit_bytes) - # A BlockSizes object carries heads_per_tile=None when the config dict omitted - # it; getattr then returns that None instead of the default, so coerce it back - # to 1 (the custom-kernel default) to keep the `heads_per_tile > 1` guards safe. - if heads_per_tile is None: - heads_per_tile = 1 - return bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes - - def _build_padding_segment_ids( query_seq_len: int, q_padded_len: int, @@ -455,32 +418,6 @@ def _tpu_flash_attention( check_rep=False, ) def wrap_flash_attention(query, key, value): - if attention_kernel == "tokamax_ring_custom": - # Ring attention backed by the custom dense splash kernel. q stays local, - # k/v rotate over the "context" axis (handled inside the ring kernel). - bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes = _extract_custom_block_sizes(flash_block_sizes) - if heads_per_tile > 1: - raise NotImplementedError("tokamax_ring_custom currently supports heads_per_tile == 1 only.") - query_local = query * LOG2E if use_base2_exp else query - query_local, kv_size, query_seq_len = _pad_data_for_flash(query_local, heads, bq) - key_local, _, key_seq_len = _pad_data_for_flash(key, heads, bkv) - value_local, _, _ = _pad_data_for_flash(value, heads, bkv) - - bsizes = custom_splash._BlockSizes(block_q=bq, block_kv=bkv, block_kv_compute=bkv_compute) - ring_kernel = tokamax_ring_attention_kernel.make_custom_ring_attention( - block_sizes=bsizes, - bkv_compute_in=bkv_compute_in, - orig_q_seq_len=query_seq_len, - orig_kv_seq_len=key_seq_len, - use_base2_exp=use_base2_exp, - use_experimental_scheduler=use_experimental_scheduler, - vmem_limit_bytes=vmem_limit_bytes, - ring_axis="context", - ) - vmapped_ring = jax.vmap(ring_kernel, in_axes=(0, 0, 0)) - attention_output = vmapped_ring(query_local, key_local, value_local) - return attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype) - uses_fused_kernel = block_sizes.use_fused_bwd_kernel block_q_sizes = ( block_sizes.block_q, @@ -646,7 +583,6 @@ def _ulysses_attention( use_custom_kernel: bool = False, use_base2_exp: bool = True, use_experimental_scheduler: bool = False, - use_fixed_m: bool = False, ) -> jax.Array: """Ulysses sequence-parallel attention. @@ -690,38 +626,36 @@ def wrap_ulysses_attention(query, key, value): value = jax.lax.all_to_all(value, axis_name=axis_name, split_axis=1, concat_axis=2, tiled=True) if use_custom_kernel: - if attention_mask is not None: - raise NotImplementedError( - "The custom dense splash kernel (use_custom_kernel) does not support attention_mask " - "(it only handles padding via orig_seq_len); got a non-None attention_mask." - ) - bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes = _extract_custom_block_sizes(flash_block_sizes) + bq = 4864 + bkv = 1024 + bkv_compute = 1024 + bkv_compute_in = 1024 + heads_per_tile = 1 + vmem_limit_bytes = None + + if flash_block_sizes is not None: + if isinstance(flash_block_sizes, dict): + bq = flash_block_sizes.get("block_q", None) or bq + bkv = flash_block_sizes.get("block_kv", None) or bkv + bkv_compute = flash_block_sizes.get("block_kv_compute", None) or bkv_compute + bkv_compute_in = flash_block_sizes.get("block_kv_compute_in", None) or bkv_compute_in + heads_per_tile = flash_block_sizes.get("heads_per_tile", None) or heads_per_tile + vmem_limit_bytes = flash_block_sizes.get("vmem_limit_bytes", None) or vmem_limit_bytes + else: + bq = getattr(flash_block_sizes, "block_q", None) or bq + bkv = getattr(flash_block_sizes, "block_kv", None) or bkv + bkv_compute = getattr(flash_block_sizes, "block_kv_compute", None) or bkv_compute + bkv_compute_in = getattr(flash_block_sizes, "block_kv_compute_in", None) or bkv_compute_in + heads_per_tile = getattr(flash_block_sizes, "heads_per_tile", None) or heads_per_tile + vmem_limit_bytes = getattr(flash_block_sizes, "vmem_limit_bytes", None) or vmem_limit_bytes if use_base2_exp: query = query * LOG2E - if use_fixed_m: - # k-smoothing (output-invariant): subtracting the per-row key mean - # forces every logit row to have mean 0, hence row-max >= 0 — the - # precondition that keeps the fixed-m Cauchy-Schwarz bound flush-free. - key = key - jnp.mean(key, axis=2, keepdims=True) - query, kv_size, query_seq_len = _pad_data_for_flash(query, heads, bq) key, _, key_seq_len = _pad_data_for_flash(key, heads, bkv) value, _, _ = _pad_data_for_flash(value, heads, bkv) - mk_arr = None - if use_fixed_m: - # Per-(local-)head Cauchy-Schwarz inputs over the (batch, seq) slice; - # padded rows have zero norm and never raise the max. mk[0] feeds the - # in-kernel per-query bound, mk[1] flags heads within the no-flush gate. - qf = query.astype(jnp.float32) - kf = key.astype(jnp.float32) - qn_max = jnp.sqrt((qf * qf).sum(-1)).max(axis=(0, 2)) # (local_heads,) - mk_h = jnp.sqrt((kf * kf).sum(-1)).max(axis=(0, 2)) # (local_heads,) - fixed_ok = (qn_max * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype(jnp.float32) - mk_arr = jnp.stack([mk_h, fixed_ok]) # (2, local_heads) - bsizes = custom_splash._BlockSizes(block_q=bq, block_kv=bkv, block_kv_compute=bkv_compute) splash_kernel = custom_splash.make_splash_mha( @@ -733,15 +667,10 @@ def wrap_ulysses_attention(query, key, value): use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, - use_fixed_m=use_fixed_m, ) - if use_fixed_m: - vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0, None)) - attention_output = vmapped_splash(query, key, value, mk_arr) - else: - vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0)) - attention_output = vmapped_splash(query, key, value) + vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0)) + attention_output = vmapped_splash(query, key, value) attention_output = jnp.swapaxes(attention_output, 2, 3) attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype) else: @@ -792,26 +721,7 @@ def wrap_ulysses_attention(query, key, value): "Warning, batch dimension should be shardable among the devices in data and fsdp" f" axis, batch dimension: {query.shape[0]}, devices_in_batch_sharding: {devices_in_batch_sharding}" ) - - # Fold the (CFG) batch into the heads axis around the Ulysses exchange. - # Each (batch, head) pair is an independent attention problem, so - # [B, H, S, D] -> [1, B*H, S, D] is mathematically identity — but it makes - # XLA compile the attention path as the batch=1 case. At batch=2 XLA - # otherwise places the size-2 batch in the tile sublanes ({3,0,1,2:T(2,128)} - # instead of T(8,128)) which quadruples the cost of every op touching the - # a2a tensors inside the scanned layers (measured 7.0 -> expected ~3.5 - # s/step at 720p 81f cp8 CFG). - batch = query.shape[0] - fold_batch = batch > 1 and (batch * num_heads) % num_shards == 0 - if fold_batch: - query = query.reshape(1, batch * num_heads, *query.shape[2:]) - key = key.reshape(1, batch * num_heads, *key.shape[2:]) - value = value.reshape(1, batch * num_heads, *value.shape[2:]) - x = wrap_ulysses_attention(query, key, value) - - if fold_batch: - x = x.reshape(batch, num_heads, *x.shape[2:]) x = x[:, :, :orig_q_seq_len, :] x = _reshape_heads_to_head_dim(x) @@ -973,183 +883,6 @@ def wrap_ulysses_ring_attention(query, key, value): return x -def _ulysses_ring_custom_attention( - query: jax.Array, - key: jax.Array, - value: jax.Array, - heads: int, - mesh: Mesh, - axis_names_q: AxisNames, - axis_names_kv: AxisNames, - flash_block_sizes: BlockSizes, - dtype: jnp.dtype = jnp.float32, - mask_padding_tokens: bool = True, - residual_checkpoint_name: str | None = None, - attention_mask: jax.Array = None, - ulysses_shards: int = -1, - use_base2_exp: bool = True, - use_experimental_scheduler: bool = False, - bidirectional: bool = False, - use_fixed_m: bool = False, -) -> jax.Array: - """Hybrid Ulysses + Ring (USP) with the CUSTOM splash kernel on main's mesh. - - Uses origin/main's explicit internal `(ring, ulysses)` mesh - (`_create_internal_ulysses_ring_mesh`, commit c104db51) instead of single-axis - collective sub-groups: the public `context` axis is reshaped with the Ulysses - axis innermost, so the Ulysses all-to-all stays INTRA-chip and the ring rotates - ACROSS chips. The per-shard attention is our custom splash kernel - (`make_custom_ring_attention`), not the tokamax_ring kernel main uses. - - 1. all-to-all over the (intra-chip) Ulysses axis: trade sequence for heads; - 2. ring (full ppermute) over the (cross-chip) ring axis, online-softmax merge; - 3. all-to-all back to restore the sequence-sharded / full-heads layout. - - U = ulysses_shards (from config); R = context // U. U=context -> pure - Ulysses, U=1 -> pure Ring (all on the same custom kernel). - """ - if attention_mask is not None: - raise NotImplementedError( - "ulysses_ring_custom does not support attention_mask (the custom splash kernels only " - "handle padding via orig_seq_len); got a non-None attention_mask." - ) - axis_name = "context" - num_context_shards = mesh.shape[axis_name] - num_ulysses_shards = ulysses_shards - if num_ulysses_shards <= 0: - raise ValueError("ulysses_ring_custom requires ulysses_shards to be set from config or command line.") - if num_context_shards % num_ulysses_shards != 0: - raise ValueError( - f"ulysses_ring_custom requires ulysses_shards to divide the context shard count, " - f"got context_shards={num_context_shards} and ulysses_shards={num_ulysses_shards}." - ) - num_ring_shards = num_context_shards // num_ulysses_shards - - query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_context_shards) - key, _ = _reshape_data_for_flash(key, heads, num_context_shards) - value, _ = _reshape_data_for_flash(value, heads, num_context_shards) - num_heads = query.shape[1] - if num_heads % num_ulysses_shards != 0: - raise ValueError(f"Ulysses+Ring requires heads divisible by U={num_ulysses_shards}, got heads={num_heads}.") - - bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes = _extract_custom_block_sizes(flash_block_sizes) - if heads_per_tile > 1: - raise NotImplementedError("ulysses_ring_custom currently supports heads_per_tile == 1 only.") - - internal_mesh = _create_internal_ulysses_ring_mesh(mesh, num_ring_shards, num_ulysses_shards) - ring_axis = INTERNAL_RING_AXIS - ulysses_axis = INTERNAL_ULYSSES_AXIS - - q_axis_names = nn.logical_to_mesh_axes(axis_names_q) - kv_axis_names = nn.logical_to_mesh_axes(axis_names_kv) - internal_q_axis_names = _replace_mesh_axis_names(q_axis_names, axis_name, (ring_axis, ulysses_axis)) - internal_kv_axis_names = _replace_mesh_axis_names(kv_axis_names, axis_name, (ring_axis, ulysses_axis)) - - @functools.partial( - jax.shard_map, - mesh=internal_mesh, - in_specs=(internal_q_axis_names, internal_kv_axis_names, internal_kv_axis_names), - out_specs=internal_q_axis_names, - check_vma=False, - ) - def wrap_ulysses_ring_attention(query, key, value): - # (1) Ulysses all-to-all over the (intra-chip) ulysses axis: heads -> sequence, - # so each device holds the full ring-chunk sequence with heads/U heads. - a2a = functools.partial(jax.lax.all_to_all, axis_name=ulysses_axis, tiled=True) - query = a2a(query, split_axis=1, concat_axis=2) - key = a2a(key, split_axis=1, concat_axis=2) - value = a2a(value, split_axis=1, concat_axis=2) - - if use_base2_exp: - query = query * LOG2E - - if use_fixed_m: - # K-smoothing precondition for fixed-m, computed PER SHARD (no ring pmean). - # A global mean would be a perfectly-uniform per-query logit shift, but the - # per-shard local mean differs from it by only O(1/sqrt(local_seq)), and the - # ring's outer online-softmax merge re-normalizes across shards anyway, so we - # drop the per-layer ring collective and accept the negligible shift error. - kbar = jnp.mean(key, axis=2, keepdims=True) - key = key - kbar - - query, kv_size, query_seq_len = _pad_data_for_flash(query, heads, bq) - key, _, key_seq_len = _pad_data_for_flash(key, heads, bkv) - value, _, _ = _pad_data_for_flash(value, heads, bkv) - - mk_arr = None - if use_fixed_m: - # Per-(local-)head Cauchy-Schwarz inputs, all LOCAL to this ring shard. The - # outer ring merge does an online softmax across shards, so each shard's - # kernel may use its own local max||k|| as the fixed-m bound for its own - # local keys -- no global ring pmax is needed for correctness. This removes - # the second per-layer ring collective. - qf = query.astype(jnp.float32) - kf = key.astype(jnp.float32) - qn_max = jnp.sqrt((qf * qf).sum(-1)).max(axis=(0, 2)) # (local_heads,) - mk_h = jnp.sqrt((kf * kf).sum(-1)).max(axis=(0, 2)) # (local_heads,) local - fixed_ok = (qn_max * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype(jnp.float32) - if os.environ.get("FIXED_M_FORCE_ALL", "0") == "1": - # PERF PROBE ONLY (unsafe): force every head onto the fixed-m fast path, - # bypassing the safety gate, to measure fixed-m's speed CEILING on the - # ring kernel. Output may be garbage; timing is still valid. - fixed_ok = jnp.ones_like(fixed_ok) - mk_arr = jnp.stack([mk_h, fixed_ok]) # (2, local_heads) - - bsizes = custom_splash._BlockSizes(block_q=bq, block_kv=bkv, block_kv_compute=bkv_compute) - if num_ring_shards == 1: - # (2a) R=1: the ring is trivial (no rotation) -> use the lighter dedicated - # splash kernel (fuse_reciprocal, no fp32 online-softmax residual windows). - # Same math as the 1-step ring, and it fits BQ=8448 where the ring kernel - # OOMs (its 3x residual windows). make_splash_mha returns [H, D, S]. - splash_kernel = custom_splash.make_splash_mha( - block_sizes=bsizes, - bkv_compute_in=bkv_compute_in, - orig_q_seq_len=query_seq_len, - orig_kv_seq_len=key_seq_len, - heads_per_tile=heads_per_tile, - use_base2_exp=use_base2_exp, - use_experimental_scheduler=use_experimental_scheduler, - vmem_limit_bytes=vmem_limit_bytes, - use_fixed_m=use_fixed_m, - ) - if use_fixed_m: - attention_output = jnp.swapaxes( - jax.vmap(splash_kernel, in_axes=(0, 0, 0, None))(query, key, value, mk_arr), 2, 3 - ) - else: - attention_output = jnp.swapaxes(jax.vmap(splash_kernel, in_axes=(0, 0, 0))(query, key, value), 2, 3) - else: - # (2b) Ring (full ppermute over the cross-chip ring axis) with the custom kernel. - # bidirectional=True -> wrap-free schedule (streams K/V both directions one hop - # at a time), for a non-wrapping ring axis. Selected by attention=ulysses_ring_custom_bidir. - ring_kernel = tokamax_ring_attention_kernel.make_custom_ring_attention( - block_sizes=bsizes, - bkv_compute_in=bkv_compute_in, - orig_q_seq_len=query_seq_len, - orig_kv_seq_len=key_seq_len, - use_base2_exp=use_base2_exp, - use_experimental_scheduler=use_experimental_scheduler, - vmem_limit_bytes=vmem_limit_bytes, - ring_axis=ring_axis, - ring_size=num_ring_shards, - bidirectional=bidirectional, - use_fixed_m=use_fixed_m, - mk=mk_arr, - ) - attention_output = jax.vmap(ring_kernel, in_axes=(0, 0, 0))(query, key, value) - attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype) - - # (3) Ulysses all-to-all back: sequence -> heads, restoring the layout. - attention_output = a2a(attention_output, split_axis=2, concat_axis=1) - return attention_output - - x = wrap_ulysses_ring_attention(query, key, value) - x = jax.lax.with_sharding_constraint(x, q_axis_names) - x = x[:, :, :orig_q_seq_len, :] - x = _reshape_heads_to_head_dim(x) - return x - - def _apply_attention_dot( query: Array, key: Array, @@ -1296,101 +1029,6 @@ def ulysses_custom_kernel(q, k, v, context): ) -@register_kernel("ulysses_ring_custom") -def ulysses_ring_custom_kernel(q, k, v, context): - return _ulysses_ring_custom_attention( - q, - k * context["scale"], - v, - context["heads"], - context["mesh"], - context["axis_names_q"], - context["axis_names_kv"], - context["flash_block_sizes"], - context["dtype"], - mask_padding_tokens=context["mask_padding_tokens"], - residual_checkpoint_name=context["residual_checkpoint_name"], - attention_mask=context["attention_mask"], - ulysses_shards=context["ulysses_shards"], - use_base2_exp=context.get("use_base2_exp", True), - use_experimental_scheduler=context.get("use_experimental_scheduler", False), - ) - - -@register_kernel("ulysses_ring_custom_fixed_m") -def ulysses_ring_custom_fixed_m_kernel(q, k, v, context): - """fixed-m variant of ulysses_ring_custom: the per-shard custom splash kernel - uses the Cauchy-Schwarz fixed-m softmax bound (no in-kernel running-max - rescale). max||k|| and the K-smoothing mean are taken LOCALLY per ring shard - (no per-layer ring collective); the outer ring online-softmax merge still - re-normalizes across shards, so per-shard bounds stay correct.""" - return _ulysses_ring_custom_attention( - q, - k * context["scale"], - v, - context["heads"], - context["mesh"], - context["axis_names_q"], - context["axis_names_kv"], - context["flash_block_sizes"], - context["dtype"], - mask_padding_tokens=context["mask_padding_tokens"], - residual_checkpoint_name=context["residual_checkpoint_name"], - attention_mask=context["attention_mask"], - ulysses_shards=context["ulysses_shards"], - use_base2_exp=context.get("use_base2_exp", True), - use_experimental_scheduler=context.get("use_experimental_scheduler", False), - use_fixed_m=True, - ) - - -@register_kernel("ulysses_ring_custom_bidir") -def ulysses_ring_custom_bidir_kernel(q, k, v, context): - """Wrap-free (bidirectional) variant of ulysses_ring_custom: the ring streams - K/V both directions one hop at a time, avoiding the diameter-length wrap hop - on a non-wrapping ring axis. Same USP split as ulysses_ring_custom otherwise.""" - return _ulysses_ring_custom_attention( - q, - k * context["scale"], - v, - context["heads"], - context["mesh"], - context["axis_names_q"], - context["axis_names_kv"], - context["flash_block_sizes"], - context["dtype"], - mask_padding_tokens=context["mask_padding_tokens"], - residual_checkpoint_name=context["residual_checkpoint_name"], - attention_mask=context["attention_mask"], - ulysses_shards=context["ulysses_shards"], - use_base2_exp=context.get("use_base2_exp", True), - use_experimental_scheduler=context.get("use_experimental_scheduler", False), - bidirectional=True, - ) - - -@register_kernel("ulysses_custom_fixed_m") -def ulysses_custom_fixed_m_kernel(q, k, v, context): - return _ulysses_attention( - q, - k * context["scale"], - v, - context["heads"], - context["mesh"], - context["axis_names_q"], - context["axis_names_kv"], - context["flash_block_sizes"], - context["dtype"], - mask_padding_tokens=context["mask_padding_tokens"], - residual_checkpoint_name=context["residual_checkpoint_name"], - attention_mask=context["attention_mask"], - use_custom_kernel=True, - use_base2_exp=context.get("use_base2_exp", True), - use_experimental_scheduler=context.get("use_experimental_scheduler", False), - use_fixed_m=True, - ) - - @register_kernel("ulysses") def ulysses_kernel(q, k, v, context): return _ulysses_attention( @@ -1490,26 +1128,6 @@ def tokamax_ring_kernel(q, k, v, context): ) -@register_kernel("tokamax_ring_custom") -def tokamax_ring_custom_kernel(q, k, v, context): - return _tpu_flash_attention( - q, - k * context["scale"], - v, - context["heads"], - context["mesh"], - context["axis_names_q"], - context["axis_names_kv"], - context["flash_block_sizes"], - context["dtype"], - attention_kernel="tokamax_ring_custom", - mask_padding_tokens=context["mask_padding_tokens"], - attention_mask=context["attention_mask"], - use_base2_exp=context.get("use_base2_exp", True), - use_experimental_scheduler=context.get("use_experimental_scheduler", False), - ) - - @register_kernel("cudnn_flash_te") def cudnn_flash_te_kernel(q, k, v, context): return _cudnn_flash_attention(q, k, v, context["heads"], context["mesh"], context["dpa_layer"]) @@ -1548,7 +1166,7 @@ def _apply_attention( seq_len_idx = 2 can_use_flash_attention = True - if attention_kernel in ["flash", "tokamax_flash", "ulysses", "ulysses_custom", "ulysses_custom_fixed_m", "ulysses_ring"]: + if attention_kernel in ["flash", "tokamax_flash", "ulysses", "ulysses_custom", "ulysses_ring"]: can_use_flash_attention = ( query.shape[seq_len_idx] >= flash_min_seq_length and key.shape[seq_len_idx] >= flash_min_seq_length @@ -1984,10 +1602,8 @@ def __init__( else: axis_names_q = (BATCH, CROSS_ATTN_HEAD, CROSS_ATTN_Q_LENGTH, D_KV) axis_names_kv = (BATCH, CROSS_ATTN_HEAD, CROSS_ATTN_KV_LENGTH, D_KV) - if attention_kernel in ("tokamax_ring", "tokamax_ring_custom", "ulysses_ring") and not is_self_attention: - attention_kernel = "tokamax_flash" # do not use ring attention for cross attention - if attention_kernel in ("ulysses_ring_custom", "ulysses_ring_custom_bidir") and not is_self_attention: - attention_kernel = "ulysses_custom" # plain ulysses (no ring) for cross attention + if attention_kernel in ("tokamax_ring", "ulysses_ring") and not is_self_attention: + attention_kernel = "tokamax_flash" self.added_kv_proj_dim = added_kv_proj_dim # New for I2V self.image_seq_len = image_seq_len # New for I2V tpu_type = get_tpu_type() @@ -2418,8 +2034,6 @@ class FlaxFluxAttention(nn.Module): out_axis_names: AxisNames = (BATCH, LENGTH, EMBED) precision: jax.lax.Precision = None qkv_bias: bool = False - use_base2_exp: bool = False - use_experimental_scheduler: bool = False def setup(self): if self.attention_kernel in {"flash", "cudnn_flash_te"} and self.mesh is None: @@ -2439,8 +2053,6 @@ def setup(self): flash_block_sizes=self.flash_block_sizes, dtype=self.dtype, float32_qk_product=False, - use_base2_exp=self.use_base2_exp, - use_experimental_scheduler=self.use_experimental_scheduler, ) kernel_axes = ("embed", "heads") @@ -2521,60 +2133,42 @@ def __call__( attention_mask=None, image_rotary_emb=None, ): - B, L = hidden_states.shape[:2] - # Deduce dimensions cleanly from class attributes - H, D = self.heads, self.dim_head - qkv_proj = self.qkv(hidden_states) - qkv_proj = checkpoint_name(qkv_proj, "img_qkv_proj") - - qkv_proj = qkv_proj.reshape(B, L, 3, H, D) - query_proj, key_proj, value_proj = jnp.split(qkv_proj, 3, axis=2) - query_proj = query_proj.squeeze(2) - key_proj = key_proj.squeeze(2) - value_proj = value_proj.squeeze(2) + B, L = hidden_states.shape[:2] + H, D, K = self.heads, qkv_proj.shape[-1] // (self.heads * 3), 3 + qkv_proj = qkv_proj.reshape(B, L, K, H, D).transpose(2, 0, 3, 1, 4) + query_proj, key_proj, value_proj = qkv_proj query_proj = self.query_norm(query_proj) + key_proj = self.key_norm(key_proj) if encoder_hidden_states is not None: - B_enc, L_txt = encoder_hidden_states.shape[:2] encoder_qkv_proj = self.encoder_qkv(encoder_hidden_states) - encoder_qkv_proj = checkpoint_name(encoder_qkv_proj, "txt_qkv_proj") - encoder_qkv_proj = encoder_qkv_proj.reshape(B_enc, L_txt, 3, H, D) - enc_query_proj, enc_key_proj, enc_value_proj = jnp.split(encoder_qkv_proj, 3, axis=2) - enc_query_proj = enc_query_proj.squeeze(2) - enc_key_proj = enc_key_proj.squeeze(2) - enc_value_proj = enc_value_proj.squeeze(2) - - encoder_query_proj = self.encoder_query_norm(enc_query_proj) - encoder_key_proj = self.encoder_key_norm(enc_key_proj) + B, L = encoder_hidden_states.shape[:2] + H, D, K = self.heads, encoder_qkv_proj.shape[-1] // (self.heads * 3), 3 + encoder_qkv_proj = encoder_qkv_proj.reshape(B, L, K, H, D).transpose(2, 0, 3, 1, 4) + encoder_query_proj, encoder_key_proj, encoder_value_proj = encoder_qkv_proj - query_proj = jnp.concatenate((encoder_query_proj, query_proj), axis=1) - key_proj = jnp.concatenate((encoder_key_proj, key_proj), axis=1) - value_proj = jnp.concatenate((enc_value_proj, value_proj), axis=1) + encoder_query_proj = self.encoder_query_norm(encoder_query_proj) - # query_proj = nn.with_logical_constraint(query_proj, self.query_axis_names) - # key_proj = nn.with_logical_constraint(key_proj, self.key_axis_names) - # value_proj = nn.with_logical_constraint(value_proj, self.value_axis_names) + encoder_key_proj = self.encoder_key_norm(encoder_key_proj) - image_rotary_emb = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2) - - query_proj = query_proj.swapaxes(1, 2) - key_proj = key_proj.swapaxes(1, 2) - query_proj, key_proj = apply_rope(query_proj, key_proj, image_rotary_emb) - query_proj = query_proj.swapaxes(1, 2) - key_proj = key_proj.swapaxes(1, 2) + query_proj = jnp.concatenate((encoder_query_proj, query_proj), axis=2) + key_proj = jnp.concatenate((encoder_key_proj, key_proj), axis=2) + value_proj = jnp.concatenate((encoder_value_proj, value_proj), axis=2) - query_proj = query_proj.reshape(B, -1, H * D) - key_proj = key_proj.reshape(B, -1, H * D) - value_proj = value_proj.reshape(B, -1, H * D) - - if encoder_hidden_states is not None: query_proj = nn.with_logical_constraint(query_proj, self.query_axis_names) key_proj = nn.with_logical_constraint(key_proj, self.key_axis_names) value_proj = nn.with_logical_constraint(value_proj, self.value_axis_names) + image_rotary_emb = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2) + query_proj, key_proj = apply_rope(query_proj, key_proj, image_rotary_emb) + + query_proj = query_proj.transpose(0, 2, 1, 3).reshape(query_proj.shape[0], query_proj.shape[2], -1) + key_proj = key_proj.transpose(0, 2, 1, 3).reshape(key_proj.shape[0], key_proj.shape[2], -1) + value_proj = value_proj.transpose(0, 2, 1, 3).reshape(value_proj.shape[0], value_proj.shape[2], -1) + attn_output = self.attention_op.apply_attention(query_proj, key_proj, value_proj, attention_mask=attention_mask) context_attn_output = None diff --git a/src/maxdiffusion/models/embeddings_flax.py b/src/maxdiffusion/models/embeddings_flax.py index 772ee8122..30a1cef45 100644 --- a/src/maxdiffusion/models/embeddings_flax.py +++ b/src/maxdiffusion/models/embeddings_flax.py @@ -446,7 +446,7 @@ def __call__(self, ids): pos = ids.astype(self.dtype) freqs_dtype = self.dtype for i in range(n_axes): - out = get_1d_rotary_pos_embed(self.axes_dim[i], pos[..., i], freqs_dtype=freqs_dtype) + out = get_1d_rotary_pos_embed(self.axes_dim[i], pos[..., i], theta=self.theta, freqs_dtype=freqs_dtype) out_freqs.append(out) out_freqs = jnp.concatenate(out_freqs, axis=1) diff --git a/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py b/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py index 266896b99..df21fa8b8 100644 --- a/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py +++ b/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py @@ -29,8 +29,6 @@ from .... import common_types from ....common_types import BlockSizes from ....utils import BaseOutput -from ...gradient_checkpoint import GradientCheckpointType, SKIP_GRADIENT_CHECKPOINT_KEY -from jax.ad_checkpoint import checkpoint_name AxisNames = common_types.AxisNames BATCH = common_types.BATCH @@ -52,42 +50,41 @@ class Transformer2DModelOutput(BaseOutput): sample: jnp.ndarray -class MlpAndOutputBlock(nn.Module): +class FlaxSwiGLUFeedForward(nn.Module): dim: int - mlp_ratio: float = 4.0 + dim_out: int + mult: float = 3.0 dtype: jnp.dtype = jnp.float32 weights_dtype: jnp.dtype = jnp.float32 precision: jax.lax.Precision = None def setup(self): - self.mlp_hidden_dim = int(self.dim * self.mlp_ratio) - self.lin_mlp = nn.Dense( - self.mlp_hidden_dim, + inner_dim = int(self.dim * self.mult) + self.linear_in = nn.Dense( + inner_dim * 2, + use_bias=False, kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")), - bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)), dtype=self.dtype, param_dtype=self.weights_dtype, precision=self.precision, + name="linear_in", ) - self.mlp_act = nn.gelu - self.linear2 = nn.Dense( - self.dim, + self.linear_out = nn.Dense( + self.dim_out, + use_bias=False, kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")), - bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)), dtype=self.dtype, param_dtype=self.weights_dtype, precision=self.precision, + name="linear_out", ) - def __call__(self, x, attn_output, gate, residual): - mlp = self.lin_mlp(x) - attn_mlp = jnp.concatenate([attn_output, self.mlp_act(mlp)], axis=2) - attn_mlp = nn.with_logical_constraint(attn_mlp, ("activation_batch", None, "mlp")) - hidden_states = self.linear2(attn_mlp) - hidden_states = checkpoint_name(hidden_states, "lin2_hidden_states") - hidden_states = gate * hidden_states - hidden_states = residual + hidden_states - return hidden_states + def __call__(self, x): + x = self.linear_in(x) + x1, x2 = jnp.split(x, 2, axis=-1) + x = nn.silu(x1) * x2 + x = self.linear_out(x) + return x class FluxSingleTransformerBlock(nn.Module): @@ -115,18 +112,29 @@ class FluxSingleTransformerBlock(nn.Module): dtype: jnp.dtype = jnp.float32 weights_dtype: jnp.dtype = jnp.float32 precision: jax.lax.Precision = None - use_base2_exp: bool = False - use_experimental_scheduler: bool = False + use_global_modulation: bool = False # Added flag! + use_swiglu: bool = False # Added flag! def setup(self): self.mlp_hidden_dim = int(self.dim * self.mlp_ratio) - self.norm = AdaLayerNormZeroSingle( - self.dim, dtype=self.dtype, weights_dtype=self.weights_dtype, precision=self.precision - ) - - self.lin_qkv = nn.Dense( - self.dim * 3, + if self.use_global_modulation: + self.norm = nn.LayerNorm( + use_bias=False, + use_scale=False, + epsilon=1e-6, + dtype=self.dtype, + param_dtype=self.weights_dtype, + ) + else: + self.norm = AdaLayerNormZeroSingle( + self.dim, dtype=self.dtype, weights_dtype=self.weights_dtype, precision=self.precision + ) + + out_dim = self.dim * 3 + (2 * self.mlp_hidden_dim if self.use_swiglu else self.mlp_hidden_dim) + self.linear1 = nn.Dense( + out_dim, + use_bias=not self.use_swiglu, kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")), bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)), dtype=self.dtype, @@ -134,14 +142,16 @@ def setup(self): precision=self.precision, ) - self.mlp_and_out = nn.remat(MlpAndOutputBlock, prevent_cse=True)( - dim=self.dim, - mlp_ratio=self.mlp_ratio, + self.mlp_act = nn.gelu + self.linear2 = nn.Dense( + self.dim, + use_bias=not self.use_swiglu, + kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")), + bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)), dtype=self.dtype, - weights_dtype=self.weights_dtype, + param_dtype=self.weights_dtype, precision=self.precision, ) - self.attn = FlaxFluxAttention( query_dim=self.dim, heads=self.num_attention_heads, @@ -151,35 +161,36 @@ def setup(self): attention_kernel=self.attention_kernel, mesh=self.mesh, flash_block_sizes=self.flash_block_sizes, - use_base2_exp=self.use_base2_exp, - use_experimental_scheduler=self.use_experimental_scheduler, ) - def __call__(self, hidden_states, temb, image_rotary_emb=None): + def __call__(self, hidden_states, temb=None, image_rotary_emb=None, temb_mod=None): residual = hidden_states - - # FIX: Constrain inputs using valid config parameters (None skips sequence length axis parsing) - hidden_states = nn.with_logical_constraint(hidden_states, ("activation_batch", None, "mlp")) - - norm_hidden_states, gate = self.norm(hidden_states, emb=temb) - - qkv = self.lin_qkv(norm_hidden_states) - qkv = checkpoint_name(qkv, "lin1_norm_hidden_states") - qkv = nn.with_logical_constraint(qkv, ("activation_batch", None, "mlp")) + if self.use_global_modulation: + shift_msa, scale_msa, gate = jnp.split(temb_mod, 3, axis=-1) + # Unsqueeze sequence dimension for broadcasting when batch_size > 1 + shift_msa = jnp.expand_dims(shift_msa, axis=1) + scale_msa = jnp.expand_dims(scale_msa, axis=1) + gate = jnp.expand_dims(gate, axis=1) + + norm_hidden_states = self.norm(hidden_states) + norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa + else: + norm_hidden_states, gate = self.norm(hidden_states, emb=temb) + qkv, mlp = jnp.split(self.linear1(norm_hidden_states), [3 * self.dim], axis=-1) + mlp = nn.with_logical_constraint(mlp, ("activation_batch", "activation_length", "activation_embed")) + qkv = nn.with_logical_constraint(qkv, ("activation_batch", "activation_length", "activation_embed")) B, L = hidden_states.shape[:2] H, D, K = self.num_attention_heads, qkv.shape[-1] // (self.num_attention_heads * 3), 3 - - qkv_proj = qkv.reshape(B, L, K, H, D) - q, k, v = jnp.split(qkv_proj, 3, axis=2) - q = q.squeeze(2).swapaxes(1, 2) - k = k.squeeze(2).swapaxes(1, 2) - v = v.squeeze(2).swapaxes(1, 2) + qkv_proj = qkv.reshape(B, L, K, H, D).transpose(2, 0, 3, 1, 4) + q, k, v = qkv_proj q = self.attn.query_norm(q) k = self.attn.key_norm(k) if image_rotary_emb is not None: + # since this function returns image_rotary_emb and passes it between layers, + # we do not want to modify it image_rotary_emb_reordered = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2) q, k = apply_rope(q, k, image_rotary_emb_reordered) @@ -188,10 +199,18 @@ def __call__(self, hidden_states, temb, image_rotary_emb=None): v = v.transpose(0, 2, 1, 3).reshape(v.shape[0], v.shape[2], -1) attn_output = self.attn.attention_op.apply_attention(q, k, v) - attn_output = checkpoint_name(attn_output, "attn_output") - hidden_states = self.mlp_and_out(norm_hidden_states, attn_output, gate, residual) + if self.use_swiglu: + mlp1, mlp2 = jnp.split(mlp, 2, axis=-1) + mlp_activated = nn.silu(mlp1) * mlp2 + else: + mlp_activated = self.mlp_act(mlp) + attn_mlp = jnp.concatenate([attn_output, mlp_activated], axis=2) + attn_mlp = nn.with_logical_constraint(attn_mlp, ("activation_batch", "activation_length", "activation_embed")) + hidden_states = self.linear2(attn_mlp) + hidden_states = gate * hidden_states + hidden_states = residual + hidden_states if hidden_states.dtype == jnp.float16: hidden_states = jnp.clip(hidden_states, -65504, 65504) @@ -211,11 +230,12 @@ class FluxTransformerBlock(nn.Module): context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the processing of `context` conditions. """ + dim: int num_attention_heads: int attention_head_dim: int qk_norm: str = "rms_norm" - eps: float = 1e-6 + eps: int = 1e-6 flash_min_seq_length: int = 4096 flash_block_sizes: BlockSizes = None mesh: jax.sharding.Mesh = None @@ -225,13 +245,24 @@ class FluxTransformerBlock(nn.Module): mlp_ratio: float = 4.0 qkv_bias: bool = False attention_kernel: str = "dot_product" - use_base2_exp: bool = False - use_experimental_scheduler: bool = False + use_global_modulation: bool = False # Added flag! + use_swiglu: bool = False # Added flag! def setup(self): - # These contain the parameter projections ("lin"), optimize them using your updated AdaLayerNorm class - self.img_norm1 = AdaLayerNormZero(self.dim, dtype=self.dtype, weights_dtype=self.weights_dtype, precision=self.precision) - self.txt_norm1 = AdaLayerNormZero(self.dim, dtype=self.dtype, weights_dtype=self.weights_dtype, precision=self.precision) + if self.use_global_modulation: + self.img_norm1 = nn.LayerNorm( + use_bias=False, use_scale=False, epsilon=self.eps, dtype=self.dtype, param_dtype=self.weights_dtype + ) + self.txt_norm1 = nn.LayerNorm( + use_bias=False, use_scale=False, epsilon=self.eps, dtype=self.dtype, param_dtype=self.weights_dtype + ) + else: + self.img_norm1 = AdaLayerNormZero( + self.dim, dtype=self.dtype, weights_dtype=self.weights_dtype, precision=self.precision + ) + self.txt_norm1 = AdaLayerNormZero( + self.dim, dtype=self.dtype, weights_dtype=self.weights_dtype, precision=self.precision + ) self.attn = FlaxFluxAttention( query_dim=self.dim, @@ -243,150 +274,188 @@ def setup(self): attention_kernel=self.attention_kernel, mesh=self.mesh, flash_block_sizes=self.flash_block_sizes, - use_base2_exp=self.use_base2_exp, - use_experimental_scheduler=self.use_experimental_scheduler, ) - # REMOVED: self.img_norm2 and self.txt_norm2 completely to stop HBM memory spilling. - # The mathematical reductions are handled natively below. - - self.img_mlp = nn.Sequential([ - nn.Dense( - int(self.dim * self.mlp_ratio), - use_bias=True, - kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")), - bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)), - dtype=self.dtype, - param_dtype=self.weights_dtype, - precision=self.precision, - ), - nn.gelu, - nn.Dense( - self.dim, - use_bias=True, - kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")), - bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)), - dtype=self.dtype, - param_dtype=self.weights_dtype, - precision=self.precision, - ), - ]) - - self.txt_mlp = nn.Sequential([ - nn.Dense( - int(self.dim * self.mlp_ratio), - use_bias=True, - kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")), - bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)), - dtype=self.dtype, - param_dtype=self.weights_dtype, - precision=self.precision, - ), - nn.gelu, - nn.Dense( - self.dim, - use_bias=True, - kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")), - bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)), - dtype=self.dtype, - param_dtype=self.weights_dtype, - precision=self.precision, - ), - ]) - - def __call__(self, hidden_states, encoder_hidden_states, temb, image_rotary_emb=None): - # Enforce active partitioning based on your FSDP setup config - hidden_states = nn.with_logical_constraint(hidden_states, ("activation_batch", None, "mlp")) - encoder_hidden_states = nn.with_logical_constraint(encoder_hidden_states, ("activation_batch", None, "mlp")) - - # 1. First Adaptive Normalization Pass - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.img_norm1(hidden_states, emb=temb) - norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.txt_norm1( - encoder_hidden_states, emb=temb + self.img_norm2 = nn.LayerNorm( + use_bias=False, + use_scale=False, + epsilon=self.eps, + dtype=self.dtype, + param_dtype=self.weights_dtype, + ) + if self.use_swiglu: + self.img_mlp = FlaxSwiGLUFeedForward( + dim=self.dim, + dim_out=self.dim, + mult=self.mlp_ratio, + dtype=self.dtype, + weights_dtype=self.weights_dtype, + precision=self.precision, + name="img_mlp", + ) + else: + self.img_mlp = nn.Sequential( + [ + nn.Dense( + int(self.dim * self.mlp_ratio), + use_bias=True, + kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")), + bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)), + dtype=self.dtype, + param_dtype=self.weights_dtype, + precision=self.precision, + ), + nn.gelu, + nn.Dense( + self.dim, + use_bias=True, + kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")), + bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)), + dtype=self.dtype, + param_dtype=self.weights_dtype, + precision=self.precision, + ), + ], + name="img_mlp", + ) + + self.txt_norm2 = nn.LayerNorm( + use_bias=False, + use_scale=False, + epsilon=self.eps, + dtype=self.dtype, + param_dtype=self.weights_dtype, ) + if self.use_swiglu: + self.txt_mlp = FlaxSwiGLUFeedForward( + dim=self.dim, + dim_out=self.dim, + mult=self.mlp_ratio, + dtype=self.dtype, + weights_dtype=self.weights_dtype, + precision=self.precision, + name="txt_mlp", + ) + else: + self.txt_mlp = nn.Sequential( + [ + nn.Dense( + int(self.dim * self.mlp_ratio), + use_bias=True, + kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")), + bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)), + dtype=self.dtype, + param_dtype=self.weights_dtype, + precision=self.precision, + ), + nn.gelu, + nn.Dense( + self.dim, + use_bias=True, + kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("mlp", "embed")), + bias_init=nn.with_logical_partitioning(nn.initializers.zeros, (None,)), + dtype=self.dtype, + param_dtype=self.weights_dtype, + precision=self.precision, + ), + ], + name="txt_mlp", + ) + + # let chunk size default to None + self._chunk_size = None + self._chunk_dim = 0 - # 2. Attention Mechanics + def __call__( + self, hidden_states, encoder_hidden_states, temb=None, image_rotary_emb=None, temb_mod_img=None, temb_mod_txt=None + ): + if self.use_global_modulation: + (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp) = jnp.split(temb_mod_img, 6, axis=-1) + (c_shift_msa, c_scale_msa, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp) = jnp.split(temb_mod_txt, 6, axis=-1) + + # Unsqueeze sequence dimension for broadcasting when batch_size > 1 + shift_msa = jnp.expand_dims(shift_msa, axis=1) + scale_msa = jnp.expand_dims(scale_msa, axis=1) + gate_msa = jnp.expand_dims(gate_msa, axis=1) + shift_mlp = jnp.expand_dims(shift_mlp, axis=1) + scale_mlp = jnp.expand_dims(scale_mlp, axis=1) + gate_mlp = jnp.expand_dims(gate_mlp, axis=1) + + c_shift_msa = jnp.expand_dims(c_shift_msa, axis=1) + c_scale_msa = jnp.expand_dims(c_scale_msa, axis=1) + c_gate_msa = jnp.expand_dims(c_gate_msa, axis=1) + c_shift_mlp = jnp.expand_dims(c_shift_mlp, axis=1) + c_scale_mlp = jnp.expand_dims(c_scale_mlp, axis=1) + c_gate_mlp = jnp.expand_dims(c_gate_mlp, axis=1) + + norm_hidden_states = self.img_norm1(hidden_states) + norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa + + norm_encoder_hidden_states = self.txt_norm1(encoder_hidden_states) + norm_encoder_hidden_states = (1 + c_scale_msa) * norm_encoder_hidden_states + c_shift_msa + else: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.img_norm1(hidden_states, emb=temb) + norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.txt_norm1( + encoder_hidden_states, emb=temb + ) + + # Attention. 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, ) - # --- IMAGE STREAM OPTIMIZATION (img_norm2) --- attn_output = gate_msa * attn_output hidden_states = hidden_states + attn_output - - # Fully fused LayerNorm + scale_mlp + shift_mlp compilation block - img_mean = jnp.mean(hidden_states, axis=-1, keepdims=True) - img_var = jnp.mean(jnp.square(hidden_states - img_mean), axis=-1, keepdims=True) - img_inv_std = jax.lax.rsqrt(img_var + self.eps) - - norm_hidden_states = (hidden_states - img_mean) * img_inv_std * (1 + scale_mlp) + shift_mlp - norm_hidden_states = nn.with_logical_constraint(norm_hidden_states, ("activation_batch", None, "mlp")) + norm_hidden_states = self.img_norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp ff_output = self.img_mlp(norm_hidden_states) - hidden_states = hidden_states + gate_mlp * ff_output + ff_output = gate_mlp * ff_output - # --- TEXT STREAM OPTIMIZATION (txt_norm2) --- + hidden_states = hidden_states + ff_output + # Process attention outputs for the `encoder_hidden_states`. context_attn_output = c_gate_msa * context_attn_output encoder_hidden_states = encoder_hidden_states + context_attn_output - # Fully fused LayerNorm + c_scale_mlp + c_shift_mlp compilation block - txt_mean = jnp.mean(encoder_hidden_states, axis=-1, keepdims=True) - txt_var = jnp.mean(jnp.square(encoder_hidden_states - txt_mean), axis=-1, keepdims=True) - txt_inv_std = jax.lax.rsqrt(txt_var + self.eps) - - norm_encoder_hidden_states = (encoder_hidden_states - txt_mean) * txt_inv_std * (1 + c_scale_mlp) + c_shift_mlp - norm_encoder_hidden_states = nn.with_logical_constraint(norm_encoder_hidden_states, ("activation_batch", None, "mlp")) + norm_encoder_hidden_states = self.txt_norm2(encoder_hidden_states) + norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp context_ff_output = self.txt_mlp(norm_encoder_hidden_states) encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output - - # Safe numerical clipping limits for half precision math execution - if encoder_hidden_states.dtype == jnp.float16 or encoder_hidden_states.dtype == jnp.bfloat16: + if encoder_hidden_states.dtype == jnp.float16: encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) - hidden_states = hidden_states.clip(-65504, 65504) - return hidden_states, encoder_hidden_states -class ScannedDoubleBlockWrapper(nn.Module): - block_kwargs: dict - - @nn.compact - def __call__(self, carry, _): - hidden_states, encoder_hidden_states, temb, image_rotary_emb = carry - - # Instantiate the pure block (no remat here) - block = FluxTransformerBlock(**self.block_kwargs) - - h_out, e_out = block( - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - temb=temb, - image_rotary_emb=image_rotary_emb, - ) - return (h_out, e_out, temb, image_rotary_emb), None - +@flax_register_to_config +class FluxTransformer2DModel(nn.Module, FlaxModelMixin, ConfigMixin): + r""" + The Transformer model introduced in Flux. -class ScannedSingleBlockWrapper(nn.Module): - block_kwargs: dict + Reference: https://blackforestlabs.ai/announcing-black-forest-labs/ - @nn.compact - def __call__(self, carry, _): - hidden_states, temb, image_rotary_emb = carry + This model inherits from [`FlaxModelMixin`]. Check the superclass documentation for it's generic methods + implemented for all models (such as downloading or saving). - # Instantiate the pure block - block = FluxSingleTransformerBlock(**self.block_kwargs) - h_out = block(hidden_states=hidden_states, temb=temb, image_rotary_emb=image_rotary_emb) - return (h_out, temb, image_rotary_emb), None + This model is also a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module) + subclass. Use it as a regular Flax Linen module and refer to the Flax documentation for all matters related to its + general usage and behavior. + Parameters: + patch_size (`int`): Patch size to turn the input data into small patches. + in_channels (`int`, *optional*, defaults to 16): The number of channels in the input. + num_layers (`int`, *optional*, defaults to 18): The number of layers of MMDiT blocks to use. + num_single_layers (`int`, *optional*, defaults to 18): The number of layers of single DiT blocks to use. + attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head. + num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention. + joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use. + pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`. + guidance_embeds (`bool`, defaults to False): Whether to use guidance embeddings. -@flax_register_to_config -class FluxTransformer2DModel(nn.Module, FlaxModelMixin, ConfigMixin): - r""" - The Transformer model introduced in Flux. """ + patch_size: int = 1 in_channels: int = 64 num_layers: int = 19 @@ -400,6 +469,7 @@ class FluxTransformer2DModel(nn.Module, FlaxModelMixin, ConfigMixin): flash_min_seq_length: int = 4096 flash_block_sizes: BlockSizes = None mesh: jax.sharding.Mesh = None + scale_shift_order: str = "shift_scale" dtype: jnp.dtype = jnp.float32 weights_dtype: jnp.dtype = jnp.float32 precision: jax.lax.Precision = None @@ -407,12 +477,12 @@ class FluxTransformer2DModel(nn.Module, FlaxModelMixin, ConfigMixin): qkv_bias: bool = True theta: int = 1000 attention_kernel: str = "dot_product" - eps: float = 1e-6 - remat_policy: str = "None" - names_which_can_be_saved: tuple = () - names_which_can_be_offloaded: tuple = () - use_base2_exp: bool = False - use_experimental_scheduler: bool = False + eps = 1e-6 + joint_attention_bias: bool = True + x_embedder_bias: bool = True + proj_out_bias: bool = True + use_global_modulation: bool = False # Added config flag! + use_swiglu: bool = False # Added config flag! def setup(self): self.out_channels = self.in_channels @@ -433,6 +503,7 @@ def setup(self): ) self.txt_in = nn.Dense( self.inner_dim, + use_bias=self.joint_attention_bias, kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), (None, "mlp")), bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("mlp",)), dtype=self.dtype, @@ -441,6 +512,7 @@ def setup(self): ) self.img_in = nn.Dense( self.inner_dim, + use_bias=self.x_embedder_bias, kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), (None, "mlp")), bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("mlp",)), dtype=self.dtype, @@ -448,85 +520,76 @@ def setup(self): precision=self.precision, ) - self.gradient_checkpoint = GradientCheckpointType.from_str(self.remat_policy) - - # 2. Apply the policy to the Module classes - # RematDoubleBlock = self.gradient_checkpoint.apply_linen(FluxTransformerBlock) - # RematSingleBlock = self.gradient_checkpoint.apply_linen(FluxSingleTransformerBlock) - - # 1. Prepare the kwargs for the double blocks - double_kwargs = { - "dim": self.inner_dim, - "num_attention_heads": self.num_attention_heads, - "attention_head_dim": self.attention_head_dim, - "attention_kernel": self.attention_kernel, - "flash_min_seq_length": self.flash_min_seq_length, - "flash_block_sizes": self.flash_block_sizes, - "mesh": self.mesh, - "dtype": self.dtype, - "weights_dtype": self.weights_dtype, - "precision": self.precision, - "mlp_ratio": self.mlp_ratio, - "qkv_bias": self.qkv_bias, - "use_base2_exp": self.use_base2_exp, - "use_experimental_scheduler": self.use_experimental_scheduler, - } - - double_policy = self.gradient_checkpoint.to_jax_policy( - names_which_can_be_saved=self.names_which_can_be_saved, - names_which_can_be_offloaded=self.names_which_can_be_offloaded, - block_type="double", - ) - - if double_policy == SKIP_GRADIENT_CHECKPOINT_KEY: - RemattedDoubleWrapper = ScannedDoubleBlockWrapper - else: - RemattedDoubleWrapper = nn.remat(ScannedDoubleBlockWrapper, prevent_cse=True, policy=double_policy) - - self.scanned_double_blocks = nn.scan( - RemattedDoubleWrapper, - variable_axes={"params": 0}, - split_rngs={"params": True, "dropout": True}, - length=self.num_layers, - metadata_params={"partition_name": None}, - )(block_kwargs=double_kwargs) - - # 3. Define pure kwargs for single blocks - single_kwargs = { - "dim": self.inner_dim, - "num_attention_heads": self.num_attention_heads, - "attention_head_dim": self.attention_head_dim, - "attention_kernel": self.attention_kernel, - "flash_min_seq_length": self.flash_min_seq_length, - "flash_block_sizes": self.flash_block_sizes, - "mesh": self.mesh, - "dtype": self.dtype, - "weights_dtype": self.weights_dtype, - "precision": self.precision, - "mlp_ratio": self.mlp_ratio, - "use_base2_exp": self.use_base2_exp, - "use_experimental_scheduler": self.use_experimental_scheduler, - } - - # 4. Force strict checkpointing on the Single Wrapper - single_policy = self.gradient_checkpoint.to_jax_policy( - names_which_can_be_saved=self.names_which_can_be_saved, - names_which_can_be_offloaded=self.names_which_can_be_offloaded, - block_type="single", - ) - - if single_policy == SKIP_GRADIENT_CHECKPOINT_KEY: - RemattedSingleWrapper = ScannedSingleBlockWrapper - else: - RemattedSingleWrapper = nn.remat(ScannedSingleBlockWrapper, prevent_cse=True, policy=single_policy) - - self.scanned_single_blocks = nn.scan( - RemattedSingleWrapper, - variable_axes={"params": 0}, - split_rngs={"params": True, "dropout": True}, - length=self.num_single_layers, - metadata_params={"partition_name": None}, - )(block_kwargs=single_kwargs) + if self.use_global_modulation: + self.double_stream_modulation_img = nn.Dense( + 6 * self.inner_dim, + use_bias=False, + kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")), + dtype=self.dtype, + param_dtype=self.weights_dtype, + precision=self.precision, + name="double_stream_modulation_img", + ) + self.double_stream_modulation_txt = nn.Dense( + 6 * self.inner_dim, + use_bias=False, + kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")), + dtype=self.dtype, + param_dtype=self.weights_dtype, + precision=self.precision, + name="double_stream_modulation_txt", + ) + self.single_stream_modulation = nn.Dense( + 3 * self.inner_dim, + use_bias=False, + kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")), + dtype=self.dtype, + param_dtype=self.weights_dtype, + precision=self.precision, + name="single_stream_modulation", + ) + + double_blocks = [] + for _ in range(self.num_layers): + double_block = FluxTransformerBlock( + dim=self.inner_dim, + num_attention_heads=self.num_attention_heads, + attention_head_dim=self.attention_head_dim, + attention_kernel=self.attention_kernel, + flash_min_seq_length=self.flash_min_seq_length, + flash_block_sizes=self.flash_block_sizes, + mesh=self.mesh, + dtype=self.dtype, + weights_dtype=self.weights_dtype, + precision=self.precision, + mlp_ratio=self.mlp_ratio, + qkv_bias=self.qkv_bias, + use_global_modulation=self.use_global_modulation, + use_swiglu=self.use_swiglu, + ) + double_blocks.append(double_block) + self.double_blocks = double_blocks + + single_blocks = [] + for _ in range(self.num_single_layers): + single_block = FluxSingleTransformerBlock( + dim=self.inner_dim, + num_attention_heads=self.num_attention_heads, + attention_head_dim=self.attention_head_dim, + attention_kernel=self.attention_kernel, + flash_min_seq_length=self.flash_min_seq_length, + flash_block_sizes=self.flash_block_sizes, + mesh=self.mesh, + dtype=self.dtype, + weights_dtype=self.weights_dtype, + precision=self.precision, + mlp_ratio=self.mlp_ratio, + use_global_modulation=self.use_global_modulation, + use_swiglu=self.use_swiglu, + ) + single_blocks.append(single_block) + + self.single_blocks = single_blocks self.norm_out = AdaLayerNormContinuous( self.inner_dim, @@ -535,6 +598,7 @@ def setup(self): dtype=self.dtype, weights_dtype=self.weights_dtype, precision=self.precision, + scale_shift_order=self.scale_shift_order, ) self.proj_out = nn.Dense( @@ -544,7 +608,7 @@ def setup(self): dtype=self.dtype, param_dtype=self.weights_dtype, precision=self.precision, - use_bias=True, + use_bias=self.proj_out_bias, ) def timestep_embedding(self, t: jax.Array, dim: int, max_period=10000, time_factor: float = 1000.0) -> jax.Array: @@ -564,9 +628,9 @@ def timestep_embedding(self, t: jax.Array, dim: int, max_period=10000, time_fact t = time_factor * t half = dim // 2 - freqs = jnp.exp(-math.log(max_period) * jnp.arange(start=0, stop=half, dtype=jnp.bfloat16) / half).astype(dtype=t.dtype) + freqs = jnp.exp(-math.log(max_period) * jnp.arange(start=0, stop=half, dtype=t.dtype) / half) - args = t[:, None].astype(jnp.bfloat16) * freqs[None] + args = t[:, None] * freqs[None] embedding = jnp.concatenate([jnp.cos(args), jnp.sin(args)], axis=-1) if dim % 2: @@ -588,13 +652,15 @@ def __call__( guidance, return_dict: bool = True, train: bool = False, + return_intermediates: bool = False, ): hidden_states = self.img_in(hidden_states) - timestep = self.timestep_embedding(timestep, 256) + timestep = self.timestep_embedding(timestep, 256, time_factor=1.0) + timestep = nn.with_logical_constraint(timestep, ("activation_batch", None)) if self.guidance_embeds: - guidance = self.timestep_embedding(guidance, 256) + guidance = self.timestep_embedding(guidance, 256, time_factor=1.0) else: guidance = None temb = ( @@ -605,6 +671,14 @@ def __call__( temb = nn.with_logical_constraint(temb, ("activation_batch", None)) + if self.use_global_modulation: + temb_silu = nn.silu(temb) + double_stream_mod_img = self.double_stream_modulation_img(temb_silu) + double_stream_mod_txt = self.double_stream_modulation_txt(temb_silu) + single_stream_mod = self.single_stream_modulation(temb_silu) + else: + double_stream_mod_img, double_stream_mod_txt, single_stream_mod = None, None, None + encoder_hidden_states = self.txt_in(encoder_hidden_states) if txt_ids.ndim == 3: txt_ids = txt_ids[0] @@ -616,23 +690,53 @@ def __call__( image_rotary_emb = self.pe_embedder(ids) image_rotary_emb = nn.with_logical_constraint(image_rotary_emb, (None, None)) - carry = (hidden_states, encoder_hidden_states, temb, image_rotary_emb) - carry, _ = self.scanned_double_blocks(carry, None) - hidden_states, encoder_hidden_states, _, _ = carry + # Initialize intermediates collection if requested + intermediates = {} + if return_intermediates: + intermediates["temb"] = temb + intermediates["global_modulation"] = (double_stream_mod_img, double_stream_mod_txt, single_stream_mod) + intermediates["double_block_inputs"] = [] + intermediates["double_block_outputs"] = [] + intermediates["single_block_outputs"] = [] + + for double_block in self.double_blocks: + if return_intermediates: + intermediates["double_block_inputs"].append((hidden_states, encoder_hidden_states)) + hidden_states, encoder_hidden_states = double_block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + temb_mod_img=double_stream_mod_img, + temb_mod_txt=double_stream_mod_txt, + ) + if return_intermediates: + intermediates["double_block_outputs"].append((hidden_states, encoder_hidden_states)) hidden_states = jnp.concatenate([encoder_hidden_states, hidden_states], axis=1) hidden_states = nn.with_logical_constraint(hidden_states, ("activation_batch", "activation_length", "activation_embed")) - # Execute the 38 Single Blocks - carry = (hidden_states, temb, image_rotary_emb) - carry, _ = self.scanned_single_blocks(carry, None) - hidden_states, _, _ = carry + for single_block in self.single_blocks: + hidden_states = single_block( + hidden_states=hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + temb_mod=single_stream_mod, + ) + if return_intermediates: + intermediates["single_block_outputs"].append(hidden_states) + + if return_intermediates: + intermediates["before_split"] = hidden_states hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...] hidden_states = self.norm_out(hidden_states, temb) output = self.proj_out(hidden_states) + if return_intermediates: + return output, intermediates + if not return_dict: return (output,) @@ -643,30 +747,26 @@ def init_weights(self, rngs, max_sequence_length, eval_only=True): resolution = 1024 num_devices = len(jax.devices()) batch_size = 1 * num_devices - in_channels = self.in_channels // 4 - joint_attention_dim = self.joint_attention_dim - pos_id_dim = 3 - pooled_projection_dim = self.pooled_projection_dim - batch_image_shape = ( batch_size, - in_channels, + 16, # 16 to match jflux.get_noise 2 * resolution // scale_factor, 2 * resolution // scale_factor, ) + # bs, encoder_input, seq_length text_shape = ( batch_size, max_sequence_length, - joint_attention_dim, + 4096, # Sequence length of text encoder, how to get this programmatically? ) text_ids_shape = ( batch_size, max_sequence_length, - pos_id_dim, + 3, # Hardcoded to match jflux.prepare ) vec_shape = ( batch_size, - pooled_projection_dim, + 768, # Sequence length of clip, how to get this programmatically? ) img = jnp.zeros(batch_image_shape, dtype=self.dtype) bs, _, h, w = img.shape @@ -678,8 +778,11 @@ def init_weights(self, rngs, max_sequence_length, eval_only=True): txt = jnp.zeros(text_shape, dtype=self.dtype) txt_ids = jnp.zeros(text_ids_shape, dtype=self.dtype) + t_vec = jnp.full(bs, 0, dtype=self.dtype) + vec = jnp.zeros(vec_shape, dtype=self.dtype) + guidance_vec = jnp.full(bs, 4.0, dtype=self.dtype) if eval_only: diff --git a/src/maxdiffusion/models/flux/util.py b/src/maxdiffusion/models/flux/util.py index 4a44bb172..4b6aedab7 100644 --- a/src/maxdiffusion/models/flux/util.py +++ b/src/maxdiffusion/models/flux/util.py @@ -154,17 +154,9 @@ def load_flow_model(name: str, eval_shapes: dict, device: str, hf_download: bool tensors[k] = torch2jax(f.get_tensor(k)) flax_state_dict = {} cpu = jax.local_devices(backend="cpu")[0] - - double_blocks_tensors = {} - single_blocks_tensors = {} - for pt_key, tensor in tensors.items(): - if pt_key.startswith("double_blocks."): - parts = pt_key.split(".") - layer_idx = int(parts[1]) - pt_key_without_idx = "double_blocks." + ".".join(parts[2:]) - renamed_pt_key = rename_key(pt_key_without_idx) - renamed_pt_key = renamed_pt_key.replace("double_blocks", "scanned_double_blocks.FluxTransformerBlock_0") + renamed_pt_key = rename_key(pt_key) + if "double_blocks" in renamed_pt_key: renamed_pt_key = renamed_pt_key.replace("img_mlp_", "img_mlp.layers_") renamed_pt_key = renamed_pt_key.replace("txt_mlp_", "txt_mlp.layers_") renamed_pt_key = renamed_pt_key.replace("img_mod", "img_norm1") @@ -176,65 +168,14 @@ def load_flow_model(name: str, eval_shapes: dict, device: str, hf_download: bool renamed_pt_key = renamed_pt_key.replace("txt_attn.proj", "attn.e_proj") renamed_pt_key = renamed_pt_key.replace("txt_attn.norm.key_norm", "attn.encoder_key_norm") renamed_pt_key = renamed_pt_key.replace("txt_attn.norm.query_norm", "attn.encoder_query_norm") - - pt_tuple_key = tuple(renamed_pt_key.split(".")) - flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, eval_shapes, scan_layers=True) - if flax_key not in double_blocks_tensors: - double_blocks_tensors[flax_key] = {} - double_blocks_tensors[flax_key][layer_idx] = flax_tensor - continue - - elif pt_key.startswith("single_blocks."): - parts = pt_key.split(".") - layer_idx = int(parts[1]) - pt_key_without_idx = "single_blocks." + ".".join(parts[2:]) - renamed_pt_key = rename_key(pt_key_without_idx) - renamed_pt_key = renamed_pt_key.replace("single_blocks", "scanned_single_blocks.FluxSingleTransformerBlock_0") - renamed_pt_key = renamed_pt_key.replace("modulation", "norm") - renamed_pt_key = renamed_pt_key.replace("norm.key_norm", "attn.key_norm") - renamed_pt_key = renamed_pt_key.replace("norm.query_norm", "attn.query_norm") - - if "linear1" in renamed_pt_key: - if tensor.ndim == 2: - qkv_tensor = tensor[:9216, :] - mlp_tensor = tensor[9216:, :] - else: - qkv_tensor = tensor[:9216] - mlp_tensor = tensor[9216:] - qkv_pt_key = renamed_pt_key.replace("linear1", "lin_qkv") - mlp_pt_key = renamed_pt_key.replace("linear1", "mlp_and_out.lin_mlp") - - flax_key_qkv, flax_tensor_qkv = rename_key_and_reshape_tensor( - tuple(qkv_pt_key.split(".")), qkv_tensor, eval_shapes, scan_layers=True - ) - flax_key_mlp, flax_tensor_mlp = rename_key_and_reshape_tensor( - tuple(mlp_pt_key.split(".")), mlp_tensor, eval_shapes, scan_layers=True - ) - - if flax_key_qkv not in single_blocks_tensors: - single_blocks_tensors[flax_key_qkv] = {} - single_blocks_tensors[flax_key_qkv][layer_idx] = flax_tensor_qkv - - if flax_key_mlp not in single_blocks_tensors: - single_blocks_tensors[flax_key_mlp] = {} - single_blocks_tensors[flax_key_mlp][layer_idx] = flax_tensor_mlp - continue - - elif "linear2" in renamed_pt_key: - renamed_pt_key = renamed_pt_key.replace("linear2", "mlp_and_out.linear2") - - pt_tuple_key = tuple(renamed_pt_key.split(".")) - flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, eval_shapes, scan_layers=True) - if flax_key not in single_blocks_tensors: - single_blocks_tensors[flax_key] = {} - single_blocks_tensors[flax_key][layer_idx] = flax_tensor - continue - - renamed_pt_key = rename_key(pt_key) - if "guidance_in" in renamed_pt_key: + elif "guidance_in" in renamed_pt_key: renamed_pt_key = renamed_pt_key.replace("guidance_in", "time_text_embed.FlaxTimestepEmbedding_1") renamed_pt_key = renamed_pt_key.replace("in_layer", "linear_1") renamed_pt_key = renamed_pt_key.replace("out_layer", "linear_2") + elif "single_blocks" in renamed_pt_key: + renamed_pt_key = renamed_pt_key.replace("modulation", "norm") + renamed_pt_key = renamed_pt_key.replace("norm.key_norm", "attn.key_norm") + renamed_pt_key = renamed_pt_key.replace("norm.query_norm", "attn.query_norm") elif "vector_in" in renamed_pt_key or "time_in" in renamed_pt_key: renamed_pt_key = renamed_pt_key.replace("vector_in", "time_text_embed.PixArtAlphaTextProjection_0") renamed_pt_key = renamed_pt_key.replace("time_in", "time_text_embed.FlaxTimestepEmbedding_0") @@ -243,25 +184,379 @@ def load_flow_model(name: str, eval_shapes: dict, device: str, hf_download: bool elif "final_layer" in renamed_pt_key: renamed_pt_key = renamed_pt_key.replace("final_layer.linear", "proj_out") renamed_pt_key = renamed_pt_key.replace("final_layer.adaLN_modulation_1", "norm_out.Dense_0") - pt_tuple_key = tuple(renamed_pt_key.split(".")) flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, eval_shapes) flax_state_dict[flax_key] = jax.device_put(jnp.asarray(flax_tensor), device=cpu) - - # Stack double blocks - for flax_key, layers in double_blocks_tensors.items(): - sorted_indices = sorted(layers.keys()) - stacked_tensor = jnp.stack([layers[i] for i in sorted_indices], axis=0) - flax_state_dict[flax_key] = jax.device_put(stacked_tensor, device=cpu) - - # Stack single blocks - for flax_key, layers in single_blocks_tensors.items(): - sorted_indices = sorted(layers.keys()) - stacked_tensor = jnp.stack([layers[i] for i in sorted_indices], axis=0) - flax_state_dict[flax_key] = jax.device_put(stacked_tensor, device=cpu) - validate_flax_state_dict(eval_shapes, flax_state_dict) flax_state_dict = unflatten_dict(flax_state_dict) del tensors jax.clear_caches() return flax_state_dict + + +# ----------------------------------------------------------------------------- +# Latent Packing & Unpacking Helpers +# ----------------------------------------------------------------------------- + + +def pack_latents(latents): + """ + Groups spatial 2x2 latent neighborhoods into a single channel dimension. + Transforms unpacked shape (batch_size, channels, height, width) + to packed sequence shape (batch_size, (height//2)*(width//2), channels*4). + """ + import numpy as np + import jax.numpy as jnp + + batch_size, channels, height, width = latents.shape + latents = np.reshape(latents, (batch_size, channels, height // 2, 2, width // 2, 2)) + latents = np.transpose(latents, (0, 2, 4, 1, 3, 5)) + latents = np.reshape(latents, (batch_size, (height // 2) * (width // 2), channels * 4)) + return jnp.array(latents) + + +def unpack_latents(latents, batch_size, num_channels_latents, height, width): + """ + Unpacks packed sequence of shape (batch_size, (height//16)*(width//16), channels*4) + back to the unpacked spatial grid shape (batch_size, channels, height//8, width//8). + """ + import numpy as np + + h_latent = height // 8 + w_latent = width // 8 + + # 1. Reshape to split spatial grid and packed channel blocks + latents = np.reshape(latents, (batch_size, h_latent // 2, w_latent // 2, num_channels_latents, 2, 2)) + # 2. Permute dimensions back to unpacked order + latents = np.transpose(latents, (0, 3, 1, 4, 2, 5)) + # 3. Flatten back to 4D unpacked latent shape + latents = np.reshape(latents, (batch_size, num_channels_latents, h_latent, w_latent)) + return latents + + +def unpack_latents_with_ids(x, x_ids, height, width): + """[B, H*W, C] -> [B, C, H, W] using coordinate IDs.""" + import jax.numpy as jnp + + batch_size, seq_len, ch = x.shape + x_list = [] + for b in range(batch_size): + data = x[b] + pos = x_ids[b] + h_ids = pos[:, 1].astype(jnp.int32) + w_ids = pos[:, 2].astype(jnp.int32) + flat_ids = h_ids * width + w_ids + out = jnp.zeros((height * width, ch), dtype=x.dtype) + out = out.at[flat_ids].set(data) + out = jnp.transpose(jnp.reshape(out, (height, width, ch)), (2, 0, 1)) + x_list.append(out) + return jnp.stack(x_list, axis=0) + + +def unpatchify_latents(latents): + """Reverses the 2x2 spatial patch grouping: [B, C, H, W] -> [B, C/4, H*2, W*2]""" + import jax.numpy as jnp + + batch_size, num_channels_latents, height, width = latents.shape + x = jnp.reshape(latents, (batch_size, num_channels_latents // 4, 2, 2, height, width)) + x = jnp.transpose(x, (0, 1, 4, 2, 5, 3)) + x = jnp.reshape(x, (batch_size, num_channels_latents // 4, height * 2, width * 2)) + return x + + +# ----------------------------------------------------------------------------- +# 4D RoPE Position Grid Helpers +# ----------------------------------------------------------------------------- + + +def prepare_latent_image_ids(batch_size, height, width): + """ + Generates positional identifiers (Height and Width coordinates) for images to build RoPE grids. + Shape: (batch_size, height * width, 4) + """ + import jax.numpy as jnp + + grid = jnp.zeros((height, width, 4), dtype=jnp.int32) + grid = grid.at[..., 1].set(jnp.arange(height)[:, None]) + grid = grid.at[..., 2].set(jnp.arange(width)[None, :]) + latent_image_ids = grid.reshape(-1, 4) + return jnp.tile(latent_image_ids[None, ...], (batch_size, 1, 1)) + + +def prepare_text_ids(batch_size, seq_len): + """ + Generates sequence index coordinate identifiers for text prompt tokens to build RoPE grids. + Shape: (batch_size, seq_len, 4) + """ + import jax.numpy as jnp + + # Text ids: [batch, seq_len, 4]. Fill text token index. + text_ids = jnp.zeros((seq_len, 4)) + # The first element is the frame index (0). The 4th is text sequence index. + text_ids = text_ids.at[..., 3].set(jnp.arange(seq_len)) + return jnp.tile(text_ids[None, ...], (batch_size, 1, 1)) + + +# ----------------------------------------------------------------------------- +# Parameter In-place Casting Helper +# ----------------------------------------------------------------------------- + + +def cast_dict_to_bfloat16_inplace(d): + """Casts a nested dictionary of JAX/numpy arrays to bfloat16 in-place, freeing memory immediately.""" + import gc + import jax.numpy as jnp + + for k, v in list(d.items()): + if isinstance(v, dict): + cast_dict_to_bfloat16_inplace(v) + elif hasattr(v, "astype"): + d[k] = jnp.array(v, dtype=jnp.bfloat16) + if hasattr(d[k], "block_until_ready"): + d[k].block_until_ready() + del v + gc.collect() + + +# ----------------------------------------------------------------------------- +# Safetensors Weight Loader & Key Converter Functions +# ----------------------------------------------------------------------------- + + +def load_and_convert_flux_klein_weights(safetensors_path, params, num_double_layers, num_single_layers): + """ + Loads PyTorch weights from safetensors and converts them to JAX parameter dictionary. + Supports dynamic layer counts (double and single stream blocks) and sharded safetensors directories. + """ + from safetensors.torch import load_file + import torch + import numpy as np + import jax.numpy as jnp + import glob + import os + import gc + + pt_state_dict = {} + if os.path.isdir(safetensors_path): + shards = glob.glob(os.path.join(safetensors_path, "*.safetensors")) + print(f"Loading sharded PyTorch weights from directory: {safetensors_path} (Found {len(shards)} shards)...") + for shard in sorted(shards): + print(f"Loading shard: {shard}...") + pt_state_dict.update(load_file(shard, device="cpu")) + else: + print(f"Loading PyTorch weights from: {safetensors_path}") + pt_state_dict = load_file(safetensors_path, device="cpu") + + print("Mapping PyTorch weights to JAX parameters...") + + first_leaf = jax.tree_util.tree_leaves(params)[0] + target_dtype = first_leaf.dtype + + def cvt(tensor, transpose=False): + if transpose: + tensor = tensor.T + return jnp.array(tensor.to(torch.float32).cpu().numpy(), dtype=target_dtype) + + # Global layers + params["txt_in"]["kernel"] = cvt(pt_state_dict.pop("context_embedder.weight"), transpose=True) + params["img_in"]["kernel"] = cvt(pt_state_dict.pop("x_embedder.weight"), transpose=True) + params["double_stream_modulation_img"]["kernel"] = cvt( + pt_state_dict.pop("double_stream_modulation_img.linear.weight"), transpose=True + ) + params["double_stream_modulation_txt"]["kernel"] = cvt( + pt_state_dict.pop("double_stream_modulation_txt.linear.weight"), transpose=True + ) + params["single_stream_modulation"]["kernel"] = cvt( + pt_state_dict.pop("single_stream_modulation.linear.weight"), transpose=True + ) + params["proj_out"]["kernel"] = cvt(pt_state_dict.pop("proj_out.weight"), transpose=True) + + # norm_out + params["norm_out"]["Dense_0"]["kernel"] = cvt(pt_state_dict.pop("norm_out.linear.weight"), transpose=True) + + # time_text_embed (Timestep Embedding) + if "time_guidance_embed.timestep_embedder.linear_1.weight" in pt_state_dict: + params["time_text_embed"]["FlaxTimestepEmbedding_0"]["linear_1"]["kernel"] = cvt( + pt_state_dict.pop("time_guidance_embed.timestep_embedder.linear_1.weight"), transpose=True + ) + if "time_guidance_embed.timestep_embedder.linear_1.bias" in pt_state_dict: + params["time_text_embed"]["FlaxTimestepEmbedding_0"]["linear_1"]["bias"] = cvt( + pt_state_dict.pop("time_guidance_embed.timestep_embedder.linear_1.bias") + ) + if "time_guidance_embed.timestep_embedder.linear_2.weight" in pt_state_dict: + params["time_text_embed"]["FlaxTimestepEmbedding_0"]["linear_2"]["kernel"] = cvt( + pt_state_dict.pop("time_guidance_embed.timestep_embedder.linear_2.weight"), transpose=True + ) + if "time_guidance_embed.timestep_embedder.linear_2.bias" in pt_state_dict: + params["time_text_embed"]["FlaxTimestepEmbedding_0"]["linear_2"]["bias"] = cvt( + pt_state_dict.pop("time_guidance_embed.timestep_embedder.linear_2.bias") + ) + + # Double Blocks + print(f"Mapping {num_double_layers} double-stream attention blocks...") + for block_idx in range(num_double_layers): + jax_db = params[f"double_blocks_{block_idx}"] + prefix = f"transformer_blocks.{block_idx}." + + # Concatenate QKV projections + to_q = pt_state_dict.pop(prefix + "attn.to_q.weight").to(torch.float32).T.cpu().numpy() + to_k = pt_state_dict.pop(prefix + "attn.to_k.weight").to(torch.float32).T.cpu().numpy() + to_v = pt_state_dict.pop(prefix + "attn.to_v.weight").to(torch.float32).T.cpu().numpy() + jax_db["attn"]["i_qkv"]["kernel"] = jnp.array(np.concatenate([to_q, to_k, to_v], axis=1), dtype=target_dtype) + + add_q = pt_state_dict.pop(prefix + "attn.add_q_proj.weight").to(torch.float32).T.cpu().numpy() + add_k = pt_state_dict.pop(prefix + "attn.add_k_proj.weight").to(torch.float32).T.cpu().numpy() + add_v = pt_state_dict.pop(prefix + "attn.add_v_proj.weight").to(torch.float32).T.cpu().numpy() + jax_db["attn"]["e_qkv"]["kernel"] = jnp.array(np.concatenate([add_q, add_k, add_v], axis=1), dtype=target_dtype) + + # Projections out + jax_db["attn"]["i_proj"]["kernel"] = cvt(pt_state_dict.pop(prefix + "attn.to_out.0.weight"), transpose=True) + jax_db["attn"]["e_proj"]["kernel"] = cvt(pt_state_dict.pop(prefix + "attn.to_add_out.weight"), transpose=True) + + # Norm scales + jax_db["attn"]["query_norm"]["scale"] = cvt(pt_state_dict.pop(prefix + "attn.norm_q.weight")) + jax_db["attn"]["key_norm"]["scale"] = cvt(pt_state_dict.pop(prefix + "attn.norm_k.weight")) + jax_db["attn"]["encoder_query_norm"]["scale"] = cvt(pt_state_dict.pop(prefix + "attn.norm_added_q.weight")) + jax_db["attn"]["encoder_key_norm"]["scale"] = cvt(pt_state_dict.pop(prefix + "attn.norm_added_k.weight")) + + # SwiGLU MLPs + jax_db["img_mlp"]["linear_in"]["kernel"] = cvt(pt_state_dict.pop(prefix + "ff.linear_in.weight"), transpose=True) + jax_db["img_mlp"]["linear_out"]["kernel"] = cvt(pt_state_dict.pop(prefix + "ff.linear_out.weight"), transpose=True) + jax_db["txt_mlp"]["linear_in"]["kernel"] = cvt(pt_state_dict.pop(prefix + "ff_context.linear_in.weight"), transpose=True) + jax_db["txt_mlp"]["linear_out"]["kernel"] = cvt( + pt_state_dict.pop(prefix + "ff_context.linear_out.weight"), transpose=True + ) + + # Single Blocks + print(f"Mapping {num_single_layers} single-stream attention blocks...") + for block_idx in range(num_single_layers): + jax_sb = params[f"single_blocks_{block_idx}"] + s_prefix = f"single_transformer_blocks.{block_idx}." + + # Joint projections + jax_sb["linear1"]["kernel"] = cvt(pt_state_dict.pop(s_prefix + "attn.to_qkv_mlp_proj.weight"), transpose=True) + jax_sb["linear2"]["kernel"] = cvt(pt_state_dict.pop(s_prefix + "attn.to_out.weight"), transpose=True) + + # Norm scales + jax_sb["attn"]["query_norm"]["scale"] = cvt(pt_state_dict.pop(s_prefix + "attn.norm_q.weight")) + jax_sb["attn"]["key_norm"]["scale"] = cvt(pt_state_dict.pop(s_prefix + "attn.norm_k.weight")) + + params = jax.tree_util.tree_map( + lambda leaf: jnp.zeros(leaf.shape, dtype=leaf.dtype) if isinstance(leaf, jax.ShapeDtypeStruct) else leaf, params + ) + del pt_state_dict + gc.collect() + print("Weight conversion complete!") + return params + + +def load_and_convert_vae_weights(safetensors_path, jax_params): + """Loads PyTorch VAE weights from safetensors, maps them to JAX, and extracts BN stats.""" + from safetensors.torch import load_file + import torch + import flax + import jax.numpy as jnp + + print(f"Loading PyTorch VAE weights from: {safetensors_path}") + pt_state_dict = load_file(safetensors_path) + + # Helper to safely convert PyTorch bfloat16 tensors to numpy float32 + def get_w(key): + return pt_state_dict[key].to(torch.float32).cpu().numpy() + + # Unfreeze JAX params so we can load the weights + jax_params = flax.core.unfreeze(jax_params) + + # Map weights + print("Mapping VAE decoder weights to JAX parameters...") + + # post_quant_conv + jax_params["post_quant_conv"]["kernel"] = jnp.array(get_w("post_quant_conv.weight").transpose(2, 3, 1, 0)) + jax_params["post_quant_conv"]["bias"] = jnp.array(get_w("post_quant_conv.bias")) + + # decoder.conv_in + jax_params["decoder"]["conv_in"]["kernel"] = jnp.array(get_w("decoder.conv_in.weight").transpose(2, 3, 1, 0)) + jax_params["decoder"]["conv_in"]["bias"] = jnp.array(get_w("decoder.conv_in.bias")) + + # decoder.mid_block + # resnets + for idx in [0, 1]: + res_jax = jax_params["decoder"]["mid_block"][f"resnets_{idx}"] + res_pt_prefix = f"decoder.mid_block.resnets.{idx}" + + res_jax["norm1"]["scale"] = jnp.array(get_w(f"{res_pt_prefix}.norm1.weight")) + res_jax["norm1"]["bias"] = jnp.array(get_w(f"{res_pt_prefix}.norm1.bias")) + res_jax["conv1"]["kernel"] = jnp.array(get_w(f"{res_pt_prefix}.conv1.weight").transpose(2, 3, 1, 0)) + res_jax["conv1"]["bias"] = jnp.array(get_w(f"{res_pt_prefix}.conv1.bias")) + + res_jax["norm2"]["scale"] = jnp.array(get_w(f"{res_pt_prefix}.norm2.weight")) + res_jax["norm2"]["bias"] = jnp.array(get_w(f"{res_pt_prefix}.norm2.bias")) + res_jax["conv2"]["kernel"] = jnp.array(get_w(f"{res_pt_prefix}.conv2.weight").transpose(2, 3, 1, 0)) + res_jax["conv2"]["bias"] = jnp.array(get_w(f"{res_pt_prefix}.conv2.bias")) + + # attentions + attn_pt_prefix = "decoder.mid_block.attentions.0" + attn_jax = jax_params["decoder"]["mid_block"]["attentions_0"] + + attn_jax["group_norm"]["scale"] = jnp.array(get_w(f"{attn_pt_prefix}.group_norm.weight")) + attn_jax["group_norm"]["bias"] = jnp.array(get_w(f"{attn_pt_prefix}.group_norm.bias")) + + attn_jax["query"]["kernel"] = jnp.array(get_w(f"{attn_pt_prefix}.to_q.weight").T) + attn_jax["query"]["bias"] = jnp.array(get_w(f"{attn_pt_prefix}.to_q.bias")) + attn_jax["key"]["kernel"] = jnp.array(get_w(f"{attn_pt_prefix}.to_k.weight").T) + attn_jax["key"]["bias"] = jnp.array(get_w(f"{attn_pt_prefix}.to_k.bias")) + attn_jax["value"]["kernel"] = jnp.array(get_w(f"{attn_pt_prefix}.to_v.weight").T) + attn_jax["value"]["bias"] = jnp.array(get_w(f"{attn_pt_prefix}.to_v.bias")) + + attn_jax["proj_attn"]["kernel"] = jnp.array(get_w(f"{attn_pt_prefix}.to_out.0.weight").T) + attn_jax["proj_attn"]["bias"] = jnp.array(get_w(f"{attn_pt_prefix}.to_out.0.bias")) + + # decoder.up_blocks + for b_idx in range(4): + up_block_jax = jax_params["decoder"][f"up_blocks_{b_idx}"] + up_block_pt = f"decoder.up_blocks.{b_idx}" + + for r_idx in range(3): + res_jax = up_block_jax[f"resnets_{r_idx}"] + res_pt = f"{up_block_pt}.resnets.{r_idx}" + + res_jax["norm1"]["scale"] = jnp.array(get_w(f"{res_pt}.norm1.weight")) + res_jax["norm1"]["bias"] = jnp.array(get_w(f"{res_pt}.norm1.bias")) + res_jax["conv1"]["kernel"] = jnp.array(get_w(f"{res_pt}.conv1.weight").transpose(2, 3, 1, 0)) + res_jax["conv1"]["bias"] = jnp.array(get_w(f"{res_pt}.conv1.bias")) + + res_jax["norm2"]["scale"] = jnp.array(get_w(f"{res_pt}.norm2.weight")) + res_jax["norm2"]["bias"] = jnp.array(get_w(f"{res_pt}.norm2.bias")) + res_jax["conv2"]["kernel"] = jnp.array(get_w(f"{res_pt}.conv2.weight").transpose(2, 3, 1, 0)) + res_jax["conv2"]["bias"] = jnp.array(get_w(f"{res_pt}.conv2.bias")) + + shortcut_key = f"{res_pt}.conv_shortcut.weight" + if shortcut_key in pt_state_dict: + res_jax["conv_shortcut"]["kernel"] = jnp.array(get_w(shortcut_key).transpose(2, 3, 1, 0)) + res_jax["conv_shortcut"]["bias"] = jnp.array(get_w(f"{res_pt}.conv_shortcut.bias")) + + if b_idx < 3: + upsampler_jax = up_block_jax["upsamplers_0"] + upsampler_pt = f"{up_block_pt}.upsamplers.0" + + upsampler_jax["conv"]["kernel"] = jnp.array(get_w(f"{upsampler_pt}.conv.weight").transpose(2, 3, 1, 0)) + upsampler_jax["conv"]["bias"] = jnp.array(get_w(f"{upsampler_pt}.conv.bias")) + + # decoder.conv_norm_out & conv_out + jax_params["decoder"]["conv_norm_out"]["scale"] = jnp.array(get_w("decoder.conv_norm_out.weight")) + jax_params["decoder"]["conv_norm_out"]["bias"] = jnp.array(get_w("decoder.conv_norm_out.bias")) + jax_params["decoder"]["conv_out"]["kernel"] = jnp.array(get_w("decoder.conv_out.weight").transpose(2, 3, 1, 0)) + jax_params["decoder"]["conv_out"]["bias"] = jnp.array(get_w("decoder.conv_out.bias")) + + # Freeze parameters + jax_params = flax.core.freeze(jax_params) + + # Extract Batch Normalization running stats + print("Extracting VAE Batch Normalization running stats...") + bn_mean = jnp.array(get_w("bn.running_mean")).reshape(1, -1, 1, 1) + bn_var = jnp.array(get_w("bn.running_var")).reshape(1, -1, 1, 1) + batch_norm_eps = 0.0001 + bn_std = jnp.sqrt(bn_var + batch_norm_eps) + + print("VAE weights and BN stats loaded successfully!") + return jax_params, bn_mean, bn_std diff --git a/src/maxdiffusion/models/normalization_flax.py b/src/maxdiffusion/models/normalization_flax.py index 5acf449c7..3d7057d6e 100644 --- a/src/maxdiffusion/models/normalization_flax.py +++ b/src/maxdiffusion/models/normalization_flax.py @@ -29,6 +29,7 @@ class AdaLayerNormContinuous(nn.Module): dtype: jnp.dtype = jnp.float32 weights_dtype: jnp.dtype = jnp.float32 precision: jax.lax.Precision = None + scale_shift_order: str = "shift_scale" @nn.compact def __call__(self, x, conditioning_embedding): @@ -42,9 +43,16 @@ def __call__(self, x, conditioning_embedding): param_dtype=self.weights_dtype, precision=self.precision, )(nn.silu(conditioning_embedding)) - shift, scale = jnp.split(emb, 2, axis=1) - shift = nn.with_logical_constraint(shift, ("activation_batch", "activation_embed")) + + if self.scale_shift_order == "scale_shift": + scale, shift = jnp.split(emb, 2, axis=1) + elif self.scale_shift_order == "shift_scale": + shift, scale = jnp.split(emb, 2, axis=1) + else: + raise ValueError(f"Unsupported scale_shift_order: {self.scale_shift_order}") + scale = nn.with_logical_constraint(scale, ("activation_batch", "activation_embed")) + shift = nn.with_logical_constraint(shift, ("activation_batch", "activation_embed")) x = nn.LayerNorm(epsilon=self.eps, use_bias=self.elementwise_affine, use_scale=self.elementwise_affine)(x) x = (1 + scale[:, None, :]) * x + shift[:, None, :] return x @@ -58,6 +66,7 @@ class AdaLayerNormZero(nn.Module): embedding_dim (`int`): The size of each embedding vector. num_embeddings (`int`): The size of the embeddings dictionary. """ + embedding_dim: int norm_type: str = "layer_norm" bias: bool = True @@ -67,10 +76,6 @@ class AdaLayerNormZero(nn.Module): @nn.compact def __call__(self, x, emb): - emb = nn.silu(emb) - - # Pretrained Flux checks: The dual block variant projects to 6 * dim - # to unpack: shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp emb = nn.Dense( 6 * self.embedding_dim, use_bias=self.bias, @@ -80,30 +85,38 @@ def __call__(self, x, emb): param_dtype=self.weights_dtype, precision=self.precision, name="lin", - )(emb) - - emb = emb[:, None, :] - - # Explicit MaxDiffusion 3D axis alignment mapping to your 'mlp' layout rule - emb = nn.with_logical_constraint(emb, ("activation_batch", None, "mlp")) - - # Slicing the 6 chunks safely within your fsdp:8, tensor:1 configuration - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = jnp.split(emb, 6, axis=-1) + )(nn.silu(emb)) + (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp) = jnp.split(emb[:, None, :], 6, axis=-1) + shift_msa = nn.with_logical_constraint(shift_msa, ("activation_batch", "activation_embed")) + scale_msa = nn.with_logical_constraint(scale_msa, ("activation_batch", "activation_embed")) + gate_msa = nn.with_logical_constraint(gate_msa, ("activation_batch", "activation_embed")) + shift_mlp = nn.with_logical_constraint(shift_mlp, ("activation_batch", "activation_embed")) + scale_mlp = nn.with_logical_constraint(scale_mlp, ("activation_batch", "activation_embed")) + gate_mlp = nn.with_logical_constraint(gate_mlp, ("activation_batch", "activation_embed")) if self.norm_type == "layer_norm": - # Fused mathematical reduction loop - mean = jnp.mean(x, axis=-1, keepdims=True) - variance = jnp.mean(jnp.square(x - mean), axis=-1, keepdims=True) - inv_std = jax.lax.rsqrt(variance + 1e-6) - - x = (x - mean) * inv_std * (1.0 + scale_msa) + shift_msa + x = nn.LayerNorm( + epsilon=1e-6, + use_bias=False, + use_scale=False, + dtype=self.dtype, + param_dtype=self.weights_dtype, + )(x) else: raise ValueError(f"Unsupported `norm_type` ({self.norm_type}) provided. Supported ones are: 'layer_norm'.") - + x = x * (1 + scale_msa) + shift_msa return x, gate_msa, shift_mlp, scale_mlp, gate_mlp class AdaLayerNormZeroSingle(nn.Module): + r""" + Norm layer adaptive layer norm zero (adaLN-Zero). + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + num_embeddings (`int`): The size of the embeddings dictionary. + """ + embedding_dim: int norm_type: str = "layer_norm" bias: bool = True @@ -114,8 +127,6 @@ class AdaLayerNormZeroSingle(nn.Module): @nn.compact def __call__(self, x, emb): emb = nn.silu(emb) - - # Matches your config layout precisely emb = nn.Dense( 3 * self.embedding_dim, use_bias=self.bias, @@ -126,27 +137,24 @@ def __call__(self, x, emb): precision=self.precision, name="lin", )(emb) - - # 1. Expand layout safely to a 3D Tensor - emb = emb[:, None, :] - - # 2. FIX: Apply verified MaxDiffusion logical rules to match the 3D footprint - # We map the channels to 'mlp' because that matches the output layout dimension of the dense layer - emb = nn.with_logical_constraint(emb, ("activation_batch", None, "mlp")) - - # 3. Slicing now happens safely within known sharding rules - shift_msa, scale_msa, gate_msa = jnp.split(emb, 3, axis=-1) - + shift_msa, scale_msa, gate_msa = jnp.split(emb[:, None, :], 3, axis=-1) + shift_msa = nn.with_logical_constraint(shift_msa, ("activation_batch", "activation_embed")) + scale_msa = nn.with_logical_constraint(scale_msa, ("activation_batch", "activation_embed")) + gate_msa = nn.with_logical_constraint(gate_msa, ("activation_batch", "activation_embed")) if self.norm_type == "layer_norm": - # Fused optimization math keeping exact pretrained weight compatibility - mean = jnp.mean(x, axis=-1, keepdims=True) - variance = jnp.mean(jnp.square(x - mean), axis=-1, keepdims=True) - inv_std = jax.lax.rsqrt(variance + 1e-6) - - x = (x - mean) * inv_std * (1.0 + scale_msa) + shift_msa + x = ( + nn.LayerNorm( + epsilon=1e-6, + use_bias=False, + use_scale=False, + dtype=self.dtype, + param_dtype=self.weights_dtype, + )(x) + * (1 + scale_msa) + + shift_msa + ) else: raise ValueError(f"Unsupported `norm_type` ({self.norm_type}) provided. Supported ones are: 'layer_norm'.") - return x, gate_msa diff --git a/src/maxdiffusion/models/qwen3_flax.py b/src/maxdiffusion/models/qwen3_flax.py new file mode 100644 index 000000000..7b78bb572 --- /dev/null +++ b/src/maxdiffusion/models/qwen3_flax.py @@ -0,0 +1,490 @@ +# Copyright 2026 The MaxDiffusion Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from typing import Any, List, Optional, Tuple +import flax.linen as nn +import jax +import jax.numpy as jnp +import numpy as np + +# ----------------------------------------------------------------------------- +# Qwen3 Configuration +# ----------------------------------------------------------------------------- + + +class FlaxQwen3Config: + + def __init__( + self, + vocab_size: int = 151936, + hidden_size: int = 2560, + intermediate_size: int = 9728, + num_hidden_layers: int = 36, + num_attention_heads: int = 32, + num_key_value_heads: int = 8, + head_dim: int = 128, + rms_norm_eps: float = 1e-6, + rope_theta: float = 1000000.0, + max_position_embeddings: int = 40960, + dtype=jnp.float32, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.rms_norm_eps = rms_norm_eps + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.dtype = dtype + + +# ----------------------------------------------------------------------------- +# Core Model Layers +# ----------------------------------------------------------------------------- + + +class FlaxQwen3RMSNorm(nn.Module): + dim: int + eps: float = 1e-6 + dtype: Any = jnp.float32 + + @nn.compact + def __call__(self, x): + x = jnp.asarray(x, self.dtype) + variance = jnp.mean(jnp.square(x), axis=-1, keepdims=True) + scale = self.param("weight", nn.initializers.ones, (self.dim,), self.dtype) + return x * jax.lax.rsqrt(variance + self.eps) * scale + + +class FlaxQwen3MLP(nn.Module): + config: FlaxQwen3Config + + @nn.compact + def __call__(self, x): + gate_proj = nn.Dense( + self.config.intermediate_size, + use_bias=False, + dtype=self.config.dtype, + name="gate_proj", + ) + up_proj = nn.Dense( + self.config.intermediate_size, + use_bias=False, + dtype=self.config.dtype, + name="up_proj", + ) + down_proj = nn.Dense( + self.config.hidden_size, + use_bias=False, + dtype=self.config.dtype, + name="down_proj", + ) + + return down_proj(jax.nn.silu(gate_proj(x)) * up_proj(x)) + + +# ----------------------------------------------------------------------------- +# Rotary Position Embeddings (RoPE) +# ----------------------------------------------------------------------------- + + +def precompute_qwen3_freqs_cis(head_dim: int, max_seq_len: int, theta: float = 1000000.0) -> Tuple[jnp.ndarray, jnp.ndarray]: + """ + Precomputes the cosine and sine tables for RoPE. + Matches standard Llama/Qwen half-half rotation layout. + """ + inv_freq = 1.0 / (theta ** (jnp.arange(0, head_dim, 2, dtype=jnp.float32) / head_dim)) + t = jnp.arange(max_seq_len, dtype=jnp.float32) + freqs = jnp.outer(t, inv_freq) # (max_seq_len, head_dim // 2) + + # Concatenate [freqs, freqs] to match Hugging Face's rotate_half layout + emb = jnp.concatenate([freqs, freqs], axis=-1) # (max_seq_len, head_dim) + + cos = jnp.cos(emb) + sin = jnp.sin(emb) + return cos, sin + + +def apply_qwen3_rotary_pos_emb( + q: jnp.ndarray, k: jnp.ndarray, cos: jnp.ndarray, sin: jnp.ndarray +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """ + Applies RoPE to Q and K tensors. + q shape: (batch, seq_len, num_heads, head_dim) + k shape: (batch, seq_len, num_kv_heads, head_dim) + cos, sin shape: (seq_len, head_dim) + """ + # Reshape cos/sin to (1, seq_len, 1, head_dim) for broadcasting + cos = cos[jnp.newaxis, :, jnp.newaxis, :] + sin = sin[jnp.newaxis, :, jnp.newaxis, :] + + def rotate_half(x): + half = x.shape[-1] // 2 + x1 = x[..., :half] + x2 = x[..., half:] + return jnp.concatenate([-x2, x1], axis=-1) + + q_rot = (q * cos) + (rotate_half(q) * sin) + k_rot = (k * cos) + (rotate_half(k) * sin) + return q_rot, k_rot + + +# ----------------------------------------------------------------------------- +# Self Attention (Grouped Query Attention) +# ----------------------------------------------------------------------------- + + +class FlaxQwen3Attention(nn.Module): + config: FlaxQwen3Config + + @nn.compact + def __call__( + self, + x: jnp.ndarray, + attention_mask: Optional[jnp.ndarray] = None, + cos_table: Optional[jnp.ndarray] = None, + sin_table: Optional[jnp.ndarray] = None, + ): + batch_size, seq_len, _ = x.shape + + # 1. Project Q, K, V + # Output sizes: Q: 4096, K: 1024, V: 1024 + q_proj = nn.Dense( + self.config.num_attention_heads * self.config.head_dim, + use_bias=False, + dtype=self.config.dtype, + name="q_proj", + ) + k_proj = nn.Dense( + self.config.num_key_value_heads * self.config.head_dim, + use_bias=False, + dtype=self.config.dtype, + name="k_proj", + ) + v_proj = nn.Dense( + self.config.num_key_value_heads * self.config.head_dim, + use_bias=False, + dtype=self.config.dtype, + name="v_proj", + ) + o_proj = nn.Dense( + self.config.hidden_size, + use_bias=False, + dtype=self.config.dtype, + name="o_proj", + ) + + # QK-Norm Layers (Head-wise, sharing scale weights of size head_dim = 128) + q_norm = FlaxQwen3RMSNorm( + dim=self.config.head_dim, + eps=self.config.rms_norm_eps, + dtype=self.config.dtype, + name="q_norm", + ) + k_norm = FlaxQwen3RMSNorm( + dim=self.config.head_dim, + eps=self.config.rms_norm_eps, + dtype=self.config.dtype, + name="k_norm", + ) + + q = q_proj(x) + k = k_proj(x) + v = v_proj(x) + + # 2. Reshape to heads first: (batch, seq_len, num_heads, head_dim) + q = q.reshape((batch_size, seq_len, self.config.num_attention_heads, self.config.head_dim)) + k = k.reshape((batch_size, seq_len, self.config.num_key_value_heads, self.config.head_dim)) + v = v.reshape((batch_size, seq_len, self.config.num_key_value_heads, self.config.head_dim)) + + # Apply QK-Norm head-wise (normalizes over the last axis of size 128) + q = q_norm(q) + k = k_norm(k) + + # 3. Apply RoPE + if cos_table is not None and sin_table is not None: + # Extract cos/sin for the current sequence length + cos = cos_table[:seq_len, :] + sin = sin_table[:seq_len, :] + q, k = apply_qwen3_rotary_pos_emb(q, k, cos, sin) + + # 4. Repeat KV heads to match Query heads (GQA) + gqa_ratio = self.config.num_attention_heads // self.config.num_key_value_heads + if gqa_ratio > 1: + k = jnp.repeat(k, gqa_ratio, axis=-2) + v = jnp.repeat(v, gqa_ratio, axis=-2) + + # 5. Transpose to (batch, num_heads, seq_len, head_dim) for attention + q = jnp.transpose(q, (0, 2, 1, 3)) + k = jnp.transpose(k, (0, 2, 1, 3)) + v = jnp.transpose(v, (0, 2, 1, 3)) + + # 6. Compute attention logits + # scores: (batch, num_heads, seq_len, seq_len) + scores = jnp.matmul(q, jnp.transpose(k, (0, 1, 3, 2))) / math.sqrt(self.config.head_dim) + + # 7. Apply causal attention mask + causal_mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=self.config.dtype)) + mask_value = jnp.where(causal_mask, 0.0, -1e10) + scores = scores + mask_value + + # 8. Apply padding attention mask if provided + if attention_mask is not None: + # attention_mask shape: (batch, seq_len) + # Reshape to (batch, 1, 1, seq_len) + p_mask = attention_mask[:, jnp.newaxis, jnp.newaxis, :] + p_mask_value = jnp.where(p_mask, 0.0, -1e10) + scores = scores + p_mask_value + + # 9. Softmax & Weighted Sum + probs = jax.nn.softmax(scores, axis=-1) + out = jnp.matmul(probs, v) # (batch, num_heads, seq_len, head_dim) + + # 10. Reshape back and project out: (batch, seq_len, hidden_size) + out = jnp.transpose(out, (0, 2, 1, 3)).reshape((batch_size, seq_len, -1)) + return o_proj(out) + + +# ----------------------------------------------------------------------------- +# Decoder Block Layer +# ----------------------------------------------------------------------------- + + +class FlaxQwen3DecoderLayer(nn.Module): + config: FlaxQwen3Config + + @nn.compact + def __call__( + self, + x: jnp.ndarray, + attention_mask: Optional[jnp.ndarray] = None, + cos_table: Optional[jnp.ndarray] = None, + sin_table: Optional[jnp.ndarray] = None, + ): + # input_layernorm + input_layernorm = FlaxQwen3RMSNorm( + dim=self.config.hidden_size, + eps=self.config.rms_norm_eps, + dtype=self.config.dtype, + name="input_layernorm", + ) + # self_attn + self_attn = FlaxQwen3Attention( + config=self.config, + name="self_attn", + ) + # post_attention_layernorm + post_attention_layernorm = FlaxQwen3RMSNorm( + dim=self.config.hidden_size, + eps=self.config.rms_norm_eps, + dtype=self.config.dtype, + name="post_attention_layernorm", + ) + # mlp + mlp = FlaxQwen3MLP( + config=self.config, + name="mlp", + ) + + # Self-Attention block (with residual) + attn_out = self_attn( + input_layernorm(x), + attention_mask=attention_mask, + cos_table=cos_table, + sin_table=sin_table, + ) + x = x + attn_out + + # MLP block (with residual) + mlp_out = mlp(post_attention_layernorm(x)) + x = x + mlp_out + + return x + + +# ----------------------------------------------------------------------------- +# Full Transformer Model +# ----------------------------------------------------------------------------- + + +class FlaxQwen3Model(nn.Module): + config: FlaxQwen3Config + + @nn.compact + def __call__( + self, + input_ids: jnp.ndarray, + attention_mask: Optional[jnp.ndarray] = None, + ) -> Tuple[jnp.ndarray, List[jnp.ndarray]]: + """ + Runs the full Qwen3-4B model. + Returns: + last_hidden_state: Output of the final layer (batch, seq_len, 2560) + all_hidden_states: List of activations from every layer, including token embeddings (length 37) + """ + batch_size, seq_len = input_ids.shape + + # 1. Token Embeddings + embed_tokens = nn.Embed( + num_embeddings=self.config.vocab_size, + features=self.config.hidden_size, + embedding_init=nn.initializers.normal(stddev=self.config.hidden_size**-0.5), + dtype=self.config.dtype, + name="embed_tokens", + ) + hidden_states = embed_tokens(input_ids) + + # Track all layer activations (including embedding layer) + all_hidden_states = [hidden_states] + + # 2. Precompute RoPE cos/sin tables + cos_table, sin_table = precompute_qwen3_freqs_cis( + head_dim=self.config.head_dim, + max_seq_len=self.config.max_position_embeddings, + theta=self.config.rope_theta, + ) + + # 3. Stacked Decoder Layers + for i in range(self.config.num_hidden_layers): + layer = FlaxQwen3DecoderLayer( + config=self.config, + name=f"layers_{i}", + ) + hidden_states = layer( + hidden_states, + attention_mask=attention_mask, + cos_table=cos_table, + sin_table=sin_table, + ) + all_hidden_states.append(hidden_states) + + # 4. Final RMSNorm + norm = FlaxQwen3RMSNorm( + dim=self.config.hidden_size, + eps=self.config.rms_norm_eps, + dtype=self.config.dtype, + name="norm", + ) + hidden_states = norm(hidden_states) + + # Keep all_hidden_states as the raw outputs of the layers, do not overwrite with final norm. + + return hidden_states, all_hidden_states + + +# ----------------------------------------------------------------------------- +# Weight Mapping & Conversion Utilities +# ----------------------------------------------------------------------------- + + +def load_and_convert_qwen3_weights(safetensors_path: str, jax_params: dict, config: FlaxQwen3Config) -> dict: + """ + Loads PyTorch weights from a safetensors file or directory of shards, + and converts them to our JAX parameter dictionary in-place. + """ + import glob + import os + from safetensors.torch import load_file + import torch + + torch_weights: dict = {} + if os.path.isdir(safetensors_path): + # Find all safetensors shards + shards = glob.glob(os.path.join(safetensors_path, "*.safetensors")) + print(f"Loading sharded Qwen3 weights from directory: {safetensors_path} (Found {len(shards)} shards)...") + for shard in sorted(shards): + print(f"Loading shard: {shard}...") + torch_weights.update(load_file(shard, device="cpu")) + else: + # Single file path + print(f"Loading Qwen3 weights from file: {safetensors_path}...") + torch_weights = load_file(safetensors_path, device="cpu") + print("PyTorch weights loaded successfully. Starting JAX parameter mapping...") + + # Helper to transpose and cast weight + def get_w(name: str, transpose: bool = True) -> np.ndarray: + nonlocal torch_weights + if name not in torch_weights: + raise KeyError(f"Weight '{name}' not found in PyTorch safetensors!") + t = torch_weights[name] + # Transpose linear layer weights (2D tensors) from (out, in) to (in, out) + if len(t.shape) == 2 and transpose: + t = t.T + return t.to(torch.float32).numpy() + + # Create mutable copy of JAX params to populate + import flax + + flat_params = flax.traverse_util.flatten_dict(jax_params) + converted_flat = {} + + for k, v in flat_params.items(): + # Reconstruct path string for debugging/matching + path_str = ".".join(k) + + # 1. Token Embeddings + if k[0] == "embed_tokens" and k[1] == "embedding": + converted_flat[k] = get_w("model.embed_tokens.weight", transpose=False) + + # 2. Decoder Layer Normalizations (RMSNorm) + elif "input_layernorm" in path_str and k[-1] == "weight": + layer_idx = k[0].split("_")[1] + converted_flat[k] = get_w(f"model.layers.{layer_idx}.input_layernorm.weight") + + elif "post_attention_layernorm" in path_str and k[-1] == "weight": + layer_idx = k[0].split("_")[1] + converted_flat[k] = get_w(f"model.layers.{layer_idx}.post_attention_layernorm.weight") + + # 3. Attention Projections & QK-Norm + elif "self_attn" in path_str and k[-1] == "kernel": + layer_idx = k[0].split("_")[1] + proj_name = k[2] # q_proj, k_proj, v_proj, o_proj + converted_flat[k] = get_w(f"model.layers.{layer_idx}.self_attn.{proj_name}.weight") + + elif "self_attn" in path_str and "q_norm" in path_str and k[-1] == "weight": + layer_idx = k[0].split("_")[1] + converted_flat[k] = get_w(f"model.layers.{layer_idx}.self_attn.q_norm.weight") + + elif "self_attn" in path_str and "k_norm" in path_str and k[-1] == "weight": + layer_idx = k[0].split("_")[1] + converted_flat[k] = get_w(f"model.layers.{layer_idx}.self_attn.k_norm.weight") + + # 4. MLP Block + elif "mlp" in path_str and k[-1] == "kernel": + layer_idx = k[0].split("_")[1] + proj_name = k[2] # gate_proj, up_proj, down_proj + converted_flat[k] = get_w(f"model.layers.{layer_idx}.mlp.{proj_name}.weight") + + # 5. Final RMSNorm + elif k[0] == "norm" and k[1] == "weight": + converted_flat[k] = get_w("model.norm.weight") + + else: + print(f"WARNING: JAX parameter '{path_str}' did not match any PyTorch weights!") + converted_flat[k] = np.zeros(v.shape, dtype=np.float32) if hasattr(v, "shape") and not isinstance(v, np.ndarray) else v + + # Clean up PyTorch memory immediately + del torch_weights + import gc + + gc.collect() + + res = flax.traverse_util.unflatten_dict(converted_flat) + return jax.tree_util.tree_map( + lambda leaf: jnp.zeros(leaf.shape, dtype=leaf.dtype) if isinstance(leaf, jax.ShapeDtypeStruct) else leaf, res + ) diff --git a/src/maxdiffusion/schedulers/scheduling_flow_match_flax.py b/src/maxdiffusion/schedulers/scheduling_flow_match_flax.py index bf88e8774..817a85f1e 100644 --- a/src/maxdiffusion/schedulers/scheduling_flow_match_flax.py +++ b/src/maxdiffusion/schedulers/scheduling_flow_match_flax.py @@ -93,6 +93,8 @@ def __init__( inverse_timesteps: bool = False, extra_one_step: bool = False, reverse_sigmas: bool = False, + use_dynamic_shifting: bool = False, + time_shift_type: str = "linear", dtype: jnp.dtype = jnp.float32, ): self.dtype = dtype @@ -142,7 +144,13 @@ def set_timesteps( if self.config.inverse_timesteps: sigmas = jnp.flip(sigmas, dims=[0]) - sigmas = current_shift * sigmas / (1 + (current_shift - 1) * sigmas) + if getattr(self.config, "use_dynamic_shifting", False): + if getattr(self.config, "time_shift_type", "exponential") == "exponential": + sigmas = jnp.exp(current_shift) / (jnp.exp(current_shift) + (1 / jnp.clip(sigmas, 1e-7, 1.0) - 1)) + else: + sigmas = current_shift * sigmas / (1 + (current_shift - 1) * sigmas) + else: + sigmas = current_shift * sigmas / (1 + (current_shift - 1) * sigmas) if self.config.reverse_sigmas: sigmas = 1 - sigmas @@ -402,3 +410,21 @@ def training_weight(self, state: FlowMatchSchedulerState, timestep: jnp.ndarray) def __len__(self) -> int: return self.config.num_train_timesteps + + +def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float: + """ + Computes the empirical time shift parameter (mu) used by Flux models + to offset sigmas dynamically based on resolution sequence length. + """ + a1, b1 = 8.73809524e-05, 1.89833333 + a2, b2 = 0.00016927, 0.45666666 + if image_seq_len > 4300: + mu = a2 * image_seq_len + b2 + return float(mu) + m_200 = a2 * image_seq_len + b2 + m_10 = a1 * image_seq_len + b1 + a = (m_200 - m_10) / 190.0 + b = m_200 - 200.0 * a + mu = a * num_steps + b + return float(mu) diff --git a/src/maxdiffusion/tests/flux2klein/generate_flux2klein_test.py b/src/maxdiffusion/tests/flux2klein/generate_flux2klein_test.py new file mode 100644 index 000000000..aee7bd936 --- /dev/null +++ b/src/maxdiffusion/tests/flux2klein/generate_flux2klein_test.py @@ -0,0 +1,410 @@ +import os +import unittest +import numpy as np +import jax +import jax.numpy as jnp + +IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true" + +# ----------------------------------------------------------------------------- +# Module-level Helper Functions for Packing & Coordinate IDs +# ----------------------------------------------------------------------------- + + +def prepare_latent_image_ids(batch_size, height, width): + """Generates 4D position coordinates (T, H, W, L) for latent tensors.""" + grid = jnp.zeros((height, width, 4), dtype=jnp.int32) + grid = grid.at[..., 1].set(jnp.arange(height)[:, None]) + grid = grid.at[..., 2].set(jnp.arange(width)[None, :]) + latent_ids = grid.reshape(-1, 4) + latent_ids = jnp.expand_dims(latent_ids, axis=0) + latent_ids = jnp.repeat(latent_ids, batch_size, axis=0) + return latent_ids + + +def pack_latents(latents): + """[B, C, H, W] -> [B, H*W, C]""" + batch_size, num_channels, height, width = latents.shape + x = jnp.reshape(latents, (batch_size, num_channels, height * width)) + x = jnp.transpose(x, (0, 2, 1)) + return x + + +def unpack_latents_with_ids(x, x_ids, height, width): + """[B, H*W, C] -> [B, C, H, W] using coordinate IDs.""" + batch_size, seq_len, ch = x.shape + x_list = [] + for b in range(batch_size): + data = x[b] + pos = x_ids[b] + h_ids = pos[:, 1].astype(jnp.int32) + w_ids = pos[:, 2].astype(jnp.int32) + flat_ids = h_ids * width + w_ids + out = jnp.zeros((height * width, ch), dtype=x.dtype) + out = out.at[flat_ids].set(data) + out = jnp.transpose(jnp.reshape(out, (height, width, ch)), (2, 0, 1)) + x_list.append(out) + return jnp.stack(x_list, axis=0) + + +def prepare_text_ids(batch_size, seq_len): + """Generates 4D position coordinates for text tokens.""" + txt_ids = jnp.zeros((seq_len, 4), dtype=jnp.int32) + txt_ids = jnp.expand_dims(txt_ids, axis=0) + return jnp.repeat(txt_ids, batch_size, axis=0) + + +class GenerateFlux2KleinTest(unittest.TestCase): + + def test_generate_random_latents_shape(self): + import torch + + latents = torch.randn((2, 32, 1024 // 8, 512 // 8)) + expected_shape = (2, 32, 1024 // 8, 512 // 8) + self.assertEqual(tuple(latents.shape), expected_shape) + + def test_packing_roundtrip_parity(self): + import torch + + latents_pt = torch.randn((2, 32, 16, 16)) + latents_jax = jnp.array(latents_pt.numpy()) + + packed_jax = pack_latents(latents_jax) + ids_jax = prepare_latent_image_ids(2, 16, 16) + unpacked_jax = unpack_latents_with_ids(packed_jax, ids_jax, 16, 16) + + diff = np.abs(np.array(unpacked_jax) - latents_pt.numpy()) + max_abs = float(np.max(diff)) + print(f"\n[UNIT TEST] Latent Packing/Unpacking Roundtrip -> Max Abs Error: {max_abs:.6e}") + self.assertEqual(max_abs, 0.0) + + def test_qwen3_flax_dummy_forward(self): + from maxdiffusion.models.qwen3_flax import FlaxQwen3Model, FlaxQwen3Config + + config = FlaxQwen3Config( + vocab_size=1000, + hidden_size=256, + intermediate_size=512, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + max_position_embeddings=512, + ) + model = FlaxQwen3Model(config) + + input_ids = jnp.array([[1, 10, 20, 30, 2]]) + key = jax.random.PRNGKey(0) + params = model.init(key, input_ids=input_ids)["params"] + + output = model.apply({"params": params}, input_ids=input_ids) + hidden_states = output[0] if isinstance(output, tuple) else output.last_hidden_state + + self.assertEqual(hidden_states.shape, (1, 5, 256)) + self.assertNotEqual(float(jnp.sum(jnp.abs(hidden_states))), 0.0) + + def test_attention_math_parity(self): + """Verifies Attention mathematical parity between JAX and PyTorch.""" + import torch + from jax.sharding import Mesh + from maxdiffusion import pyconfig + from maxdiffusion.max_utils import create_device_mesh + from diffusers.models.attention_processor import Attention as PTAttention + from maxdiffusion.models.attention_flax import FlaxAttention + + pyconfig._config = None + pyconfig.initialize([ + None, + "src/maxdiffusion/configs/base_flux2klein.yml", + "run_name=flux_test", + "output_dir=/tmp/", + "jax_cache_dir=/tmp/cache_dir", + "skip_jax_distributed_system=True", + ]) + config = pyconfig.config + devices_array = create_device_mesh(config) + mesh = Mesh(devices_array[:1, :1], config.mesh_axes) + + pt_attn = PTAttention(query_dim=3072, heads=24, dim_head=128, bias=False) + pt_attn.eval() + + jax_attn = FlaxAttention( + query_dim=3072, + heads=24, + dim_head=128, + dtype=jnp.float32, + mesh=mesh, + ) + + key = jax.random.PRNGKey(0) + dummy_x = jnp.zeros((1, 128, 3072), dtype=jnp.float32) + + with mesh, jax.set_mesh(mesh): + params = jax_attn.init(key, dummy_x)["params"] + + pt_sd = pt_attn.state_dict() + params["to_q"]["kernel"] = pt_sd["to_q.weight"].T.numpy() + params["to_k"]["kernel"] = pt_sd["to_k.weight"].T.numpy() + params["to_v"]["kernel"] = pt_sd["to_v.weight"].T.numpy() + params["to_out_0"]["kernel"] = pt_sd["to_out.0.weight"].T.numpy() + + np.random.seed(42) + pt_x = torch.randn(1, 128, 3072, dtype=torch.float32) + + with torch.no_grad(): + pt_out = pt_attn(pt_x).numpy() + + jax_out = np.array(jax_attn.apply({"params": params}, jnp.array(pt_x.numpy()))) + + diff = np.abs(jax_out - pt_out) + max_abs = float(np.max(diff)) + rmse = float(np.sqrt(np.mean(diff**2))) + + print(f"\n[UNIT TEST] Attention Parity -> Max Abs Error: {max_abs:.6e}, RMSE: {rmse:.6e}") + self.assertLess(max_abs, 5e-2) + self.assertLess(rmse, 2e-2) + + def test_flowmatch_scheduler_parity(self): + """Verifies FlowMatch Euler Discrete Scheduler stepping parity against PyTorch.""" + import torch + import jax.numpy as jnp + import numpy as np + from diffusers import FlowMatchEulerDiscreteScheduler + + pt_sched = FlowMatchEulerDiscreteScheduler(shift=1.0) + pt_sched.set_timesteps(num_inference_steps=4) + + np.random.seed(42) + pt_sample = torch.randn(1, 1024, 64, dtype=torch.float32) + pt_model_output = torch.randn(1, 1024, 64, dtype=torch.float32) + + jax_sample = jnp.array(pt_sample.numpy()) + jax_model_output = jnp.array(pt_model_output.numpy()) + + sigmas = pt_sched.sigmas.numpy() + max_abs = 0.0 + + for i in range(4): + dt = sigmas[i + 1] - sigmas[i] + jax_step = jax_sample + dt * jax_model_output + + pt_step = pt_sched.step(pt_model_output, pt_sched.timesteps[i], pt_sample).prev_sample + + diff = np.abs(np.array(jax_step) - pt_step.numpy()) + max_abs = max(max_abs, float(np.max(diff))) + + jax_sample = jax_step + pt_sample = pt_step + + print(f"\n[UNIT TEST] FlowMatch Scheduler Parity (4-Step) -> Max Abs Error: {max_abs:.6e}") + self.assertEqual(max_abs, 0.0) + + def test_double_transformer_block_dummy_parity(self): + """Verifies Double Transformer Block parity under dummy weights and zero modulation.""" + import torch + import jax + import jax.numpy as jnp + import numpy as np + import flax + import flax.linen.spmd as flax_spmd + from jax.sharding import Mesh + from maxdiffusion import pyconfig + from maxdiffusion.max_utils import create_device_mesh + from diffusers.models.transformers.transformer_flux2 import Flux2TransformerBlock + from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformerBlock + + pyconfig._config = None + pyconfig.initialize([ + None, + "src/maxdiffusion/configs/base_flux2klein.yml", + "run_name=flux_test", + "output_dir=/tmp/", + "jax_cache_dir=/tmp/cache_dir", + "skip_jax_distributed_system=True", + "weights_dtype=float32", + "activations_dtype=float32", + ]) + config = pyconfig.config + devices_array = create_device_mesh(config) + mesh = Mesh(devices_array[:1, :1], config.mesh_axes) + + pt_block = Flux2TransformerBlock(dim=3072, num_attention_heads=24, attention_head_dim=128, mlp_ratio=3.0) + pt_block.eval() + + jax_block = FluxTransformerBlock( + dim=3072, + num_attention_heads=24, + attention_head_dim=128, + mlp_ratio=3.0, + use_global_modulation=True, + use_swiglu=True, + dtype=jnp.float32, + weights_dtype=jnp.float32, + mesh=mesh, + ) + + key = jax.random.PRNGKey(0) + img_dummy = jnp.zeros((1, 256, 3072), dtype=jnp.float32) + txt_dummy = jnp.zeros((1, 512, 3072), dtype=jnp.float32) + mod_dummy = jnp.zeros((1, 18432), dtype=jnp.float32) + rope_dummy = jnp.zeros((768, 64, 4), dtype=jnp.float32) + + with mesh, jax.set_mesh(mesh): + variables = jax_block.init( + key, + img_dummy, + txt_dummy, + temb=None, + image_rotary_emb=rope_dummy, + temb_mod_img=mod_dummy, + temb_mod_txt=mod_dummy, + ) + params = variables["params"] + params = jax.tree_util.tree_map( + lambda x: (x.unbox() if isinstance(x, flax_spmd.LogicallyPartitioned) else x), + params, + is_leaf=lambda k: isinstance(k, flax_spmd.LogicallyPartitioned), + ) + params = flax.core.unfreeze(params) + + np.random.seed(42) + pt_img = torch.randn(1, 256, 3072, dtype=torch.float32) + pt_txt = torch.randn(1, 512, 3072, dtype=torch.float32) + pt_mod = torch.zeros(1, 18432, dtype=torch.float32) + + pt_cos = torch.ones(768, 128, dtype=torch.float32) + pt_sin = torch.zeros(768, 128, dtype=torch.float32) + pt_rope = (pt_cos, pt_sin) + + jax_cos = np.ones((768, 64, 2), dtype=np.float32) + jax_sin = np.zeros((768, 64, 2), dtype=np.float32) + jax_rope = np.concatenate([jax_cos, jax_sin], axis=-1) + + with torch.no_grad(): + pt_txt_out, pt_img_out = pt_block( + hidden_states=pt_img, + encoder_hidden_states=pt_txt, + temb_mod_img=pt_mod, + temb_mod_txt=pt_mod, + image_rotary_emb=pt_rope, + ) + + jax_img_out, jax_txt_out = jax_block.apply( + {"params": params}, + jnp.array(pt_img.numpy()), + jnp.array(pt_txt.numpy()), + temb=None, + image_rotary_emb=jnp.array(jax_rope), + temb_mod_img=jnp.array(pt_mod.numpy()), + temb_mod_txt=jnp.array(pt_mod.numpy()), + ) + + diff_img = np.abs(np.array(jax_img_out) - pt_img_out.numpy()) + diff_txt = np.abs(np.array(jax_txt_out) - pt_txt_out.numpy()) + max_abs = float(max(np.max(diff_img), np.max(diff_txt))) + rmse = float(np.sqrt(0.5 * (np.mean(diff_img**2) + np.mean(diff_txt**2)))) + + print(f"\n[UNIT TEST] Double Block Parity -> Max Abs Error: {max_abs:.6f}, RMSE: {rmse:.6e}") + self.assertLess(max_abs, 1e-4) + self.assertLess(rmse, 1e-4) + + def test_single_transformer_block_dummy_parity(self): + """Verifies Single Transformer Block parity under dummy weights and zero modulation.""" + import torch + import jax + import jax.numpy as jnp + import numpy as np + import flax + import flax.linen.spmd as flax_spmd + from jax.sharding import Mesh + from maxdiffusion import pyconfig + from maxdiffusion.max_utils import create_device_mesh + from diffusers.models.transformers.transformer_flux2 import Flux2SingleTransformerBlock + from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxSingleTransformerBlock + + pyconfig._config = None + pyconfig.initialize([ + None, + "src/maxdiffusion/configs/base_flux2klein.yml", + "run_name=flux_test", + "output_dir=/tmp/", + "jax_cache_dir=/tmp/cache_dir", + "skip_jax_distributed_system=True", + "weights_dtype=float32", + "activations_dtype=float32", + ]) + config = pyconfig.config + devices_array = create_device_mesh(config) + mesh = Mesh(devices_array[:1, :1], config.mesh_axes) + + pt_block = Flux2SingleTransformerBlock(dim=3072, num_attention_heads=24, attention_head_dim=128) + pt_block.eval() + + jax_block = FluxSingleTransformerBlock( + dim=3072, + num_attention_heads=24, + attention_head_dim=128, + mlp_ratio=3.0, + use_global_modulation=True, + use_swiglu=True, + dtype=jnp.float32, + weights_dtype=jnp.float32, + mesh=mesh, + ) + + key = jax.random.PRNGKey(0) + x_dummy = jnp.zeros((1, 1536, 3072), dtype=jnp.float32) + mod_dummy = jnp.zeros((1, 9216), dtype=jnp.float32) + rope_dummy = jnp.zeros((1536, 64, 4), dtype=jnp.float32) + + with mesh, jax.set_mesh(mesh): + variables = jax_block.init(key, x_dummy, temb_mod=mod_dummy, image_rotary_emb=rope_dummy) + params = variables["params"] + params = jax.tree_util.tree_map( + lambda x: (x.unbox() if isinstance(x, flax_spmd.LogicallyPartitioned) else x), + params, + is_leaf=lambda k: isinstance(k, flax_spmd.LogicallyPartitioned), + ) + params = flax.core.unfreeze(params) + + pt_sd = pt_block.state_dict() + params["linear1"]["kernel"] = pt_sd["attn.to_qkv_mlp_proj.weight"].T.numpy() + params["linear2"]["kernel"] = pt_sd["attn.to_out.weight"].T.numpy() + params["attn"]["query_norm"]["scale"] = pt_sd["attn.norm_q.weight"].numpy() + params["attn"]["key_norm"]["scale"] = pt_sd["attn.norm_k.weight"].numpy() + params = flax.core.freeze(params) + + np.random.seed(42) + pt_img = torch.randn(1, 1024, 3072, dtype=torch.float32) + pt_txt = torch.randn(1, 512, 3072, dtype=torch.float32) + pt_mod = torch.zeros(1, 9216, dtype=torch.float32) + + pt_cos = torch.ones(1536, 128, dtype=torch.float32) + pt_sin = torch.zeros(1536, 128, dtype=torch.float32) + pt_rope = (pt_cos, pt_sin) + + jax_cos = np.ones((1536, 64, 2), dtype=np.float32) + jax_sin = np.zeros((1536, 64, 2), dtype=np.float32) + jax_rope = np.concatenate([jax_cos, jax_sin], axis=-1) + + with torch.no_grad(): + pt_out = pt_block(hidden_states=pt_img, encoder_hidden_states=pt_txt, temb_mod=pt_mod, image_rotary_emb=pt_rope) + + jax_input = jnp.concatenate([jnp.array(pt_txt.numpy()), jnp.array(pt_img.numpy())], axis=1) + jax_out = jax_block.apply( + {"params": params}, + jax_input, + temb_mod=jnp.array(pt_mod.numpy()), + image_rotary_emb=jnp.array(jax_rope), + ) + + diff = np.abs(np.array(jax_out) - pt_out.numpy()) + max_abs = float(np.max(diff)) + rmse = float(np.sqrt(np.mean(diff**2))) + + print(f"\n[UNIT TEST] Single Block Parity -> Max Abs Error: {max_abs:.6f}, RMSE: {rmse:.6e}") + self.assertLess(max_abs, 1e-4) + self.assertLess(rmse, 1e-4) + + +if __name__ == "__main__": + unittest.main()