Sunday, June 14, 2026: (2,3) Handling Tensor Dimension Shifts Caused by Concatenating Instance and Class Data in Prior Preservation

Sunday, June 14, 2026

Analyzing a Diffusers Issue: FLUX DreamBooth Matrix Mismatch During PR Process

The Hugging Face diffusers library is broadly divided into two parts:

• Core Library: The core engine code located in src/diffusers.
• Examples Scripts: Training codes where users run scripts like python train_dreambooth_flux.py --options... (e.g., inside the examples/dreambooth folder).

Looking at the issue, the contributor who opened it seems to suspect that the problem lies not within the diffusers core library itself, but rather with how data is processed in the train_dreambooth_flux.py training script.

Since it triggers a RuntimeError, it is not a syntax error where the code itself is wrong; it is an error that occurs during execution. The operations seem to run fine in the earlier stages but crash at the embedding layer stage. (An embedding layer is where words like "apple," which computers cannot inherently understand, are converted into specific coordinates/vectors so the AI can process them).

In PyTorch’s C++ engine, matrices are abbreviated as mat, and sequential matrix multiplications (mat1, mat2, etc.) are performed.

According to the issue title, the input data matrix mat1 is 2 x 1536, while the weight matrix mat2 is 768 x 3072. Because the rows and columns do not match, the multiplication operation fails.

However, looking closely at the numbers, the column size of mat1 is exactly double the row size of mat2. It was originally designed to be 768, but because the user enabled the --with_prior_preservation option, the training data and the regularization data became concatenated/arranged together, resulting in this doubled shape.

• **--with_prior_preservation option:* DreamBooth fine-tuning involves giving the model specific photos to train on a subject more deeply. However, if I keep training the model on photos of my own dog as a "dog," the model might end up learning that all* dogs look exactly like mine. Enabling this option prevents this by preserving the model's existing prior knowledge.
• Class Images: This refers to the data used to prevent overfitting. When the option above is enabled, the program ingests two types of data: the training data (my dog) and the regularization class images (generic dogs).

Investigating the FLUX DreamBooth LoRA Issue: Version Regression and the Modal Environment

The user was working with FLUX, the latest image generation AI model, and using LoRA for efficient DreamBooth training. The script in question was train_dreambooth_lora_flux_advanced.py—currently one of the most trendy yet computationally heavy training setups available.

At the time of writing, the latest version of the library seems to have been v0.35.1, which was throwing the error. Interestingly, downgrading to v0.31.0 resolved the issue. When the repository was updated to support FLUX, the codebase underwent a major overhaul. During this process, the developers missed updating the tensor dimension calculations when the --with_prior_preservation option is enabled.

(Missing the tensor dimension calculations means that they forgot to implement the logic to handle the doubled data size when the --with_prior_preservation flag is turned on in the new FLUX model).

=> Consequently, comparing the source code differences between the two versions will likely lead to the solution.

Training AI models requires significant financial investment, usually involving installing multi-thousand-dollar GPUs on a local machine or renting massive servers from cloud providers like AWS. However, using a serverless platform like Modal allows the GPU to spin up only during training and shut down immediately after. This way, you only pay for the exact, short duration of the compute time used. For this setup, an NVIDIA L4 GPU was utilized — a highly powerful graphics card designed for AI data centers. With its large 24GB VRAM capacity, it is capable of running the FLUX model.

The user retrained the model keeping all parameters and arguments identical (including image size, learning rate, and the specific --with_prior_preservation option we are focusing on). In other words, changing only the Diffusers version was what caused the error to appear or disappear.

=> Therefore, a fix must be implemented so that this issue is resolved in the latest version as well.

Investigating Environment Differences: Accelerate Env and Version Upgrades

accelerate is a training acceleration library created by Hugging Face. When training AI, it determines whether to use a single GPU or multiple GPUs, and which techniques to apply for saving memory. The command accelerate env is used in the terminal to print out a complete breakdown of the current environment, such as the Python version and the GPU being used. When reporting bugs to open-source repositories, users usually include this output so others can reproduce the issue.

The user ran the accelerate env command while Modal was still building the container image (meaning the server configuration phase, not the AI generating an image). Because of this, the exact GPU model name was not captured properly in the log.

=> This indicates that the command was executed during the setup phase before the server fully initialized, meaning the missing GPU name is just a timing issue, and the core problem lies within the latest codebase.

The user also provided the accelerate env output for the older version. Comparing the two setups revealed distinct differences in configuration values between the new and old versions:

• PyTorch Version: The new version uses 2.8.0+cu129, while the old version uses 2.5.1+cu124. PyTorch underwent a major upgrade. This suggests that the internal matrix multiplication methods might have changed, or tensor shape validation has become much stricter.
• Accelerate Version: The new version uses 1.10.0, while the old version uses 1.2.1. When the --with_prior_preservation option is enabled, accelerate is involved in preparing the training data alongside the regularization data and feeding them into the GPU. The upgrade to this library might have altered how data is concatenated and distributed.

Looking closer at the logs, the old version showed PyTorch version (GPU?): 2.5.1+cu124 (False), meaning it failed to recognize the PyTorch GPU version. In the new version, it showed PyTorch accelerator: N/A. In both cases, the user likely triggered the command before the environment setup was fully complete.

=> This confirms that the unrecognized GPU status is merely an issue of when the command was executed, making it highly unlikely to be a hardware or basic configuration fault.

Pinpointing the Bug: Traceback Analysis and Version Comparison

Looking at the error log from the new version:

File ".../transformer_flux.py", line 696, in forward
  else self.time_text_embed(timestep, guidance, pooled_projections)
File ".../embeddings.py", line 1614, in forward
  pooled_projections = self.text_embedder(pooled_projection)
File ".../embeddings.py", line 2207, in forward
  hidden_states = self.linear_1(caption)

The data passes through transformer_flux.py in the FLUX model and enters the time_text_embed block. Inside, pooled_projection data is passed into text_embedder within embeddings.py. The crash happens the exact moment the caption tensor data is fed into the linear_1 layer inside embeddings.py.

=> The specific code that needs investigation is around line 2200, near linear_1 in src/diffusers/models/embeddings.py.

Now, looking at the log from the old version:

Passing `txt_ids` 3d torch.Tensor is deprecated. 
Please remove the batch dimension and pass it as a 2d torch Tensor

It warns that passing txt_ids as a 3D torch.Tensor is deprecated and instructs to remove the batch dimension and pass it as a 2D torch.Tensor. (A batch dimension represents grouping a set amount of data to process at once. For example, a batch size of 3 means processing 3 data points bundled together).

=> In the older version, the dimension handling for text data was not fully optimized, so text data simply bypassed as 3D. However, in the latest version, code to remove or modify this batch dimension has been introduced into the core Diffusers engine.

During this refactoring, the developers likely updated the code with only the standard batch size of 1 in mind. They failed to anticipate that the data size would double when a user enables the --with_prior_preservation option. This is exactly why the dimension got tangled up at 1536.

The user encountered a subprocess.CalledProcessError: Command... execution error. This bug can be fully reproduced. The --with_prior_preservation option is the core root cause of this bug. While the explicit setting is --train_batch_size=1 (default batch size of 1), it internally needs to become 2 due to the regularization class images. Additionally, the --train_text_encoder_ti (Text Encoder Textual Inversion training) option is also enabled.

Analyzing the Setup Code: Git Clones, PEFT, and Cross-Model Overlap

• CLI (Command Line Interface): A method of operating a computer by typing commands into a terminal. Because troubleshooting directly inside the Modal environment was difficult, the author shared the entire server configuration code. This includes the building image (Docker image) configuration — a setup value that configures everything in one go.

Looking at the command .run_commands("git clone https://github.com/huggingface/diffusers.git ..."), it is clear that no specific version was specified. This means it pulled the absolute latest branch. Although the user mentioned being on v0.35.1, this issue likely remains unresolved in the latest master branch as well.

The configuration includes "peft>=0.11.1,<=0.17.0". The user integrated PEFT within its latest supported range to connect LoRA with Diffusers. The dimension calculation got tangled up during the embedding layer phase right when Diffusers passes data over to the PEFT library. This is precisely why the error exploded inside train_dreambooth_lora_flux_advanced.py.

Now, looking at the setup code for the older version: .run_commands("git clone -b v0.31.0 https://github.com/huggingface/diffusers.git ...")

In the older setup, a specific version was explicitly designated for download. At that time, train_dreambooth_lora_flux_advanced.py either did not exist at all, or its structure was completely different. Following that version, the developers newly added the advanced script to complexly process the text encoder and timestep embedding for better FLUX performance. During that development process, they highly likely missed handling the doubled data size when the --with_prior_preservation option is enabled.

• Text Encoder: A translator that converts text into vectors that the AI can understand. It tokenizes the characters we type into the terminal or prompt and embeds them into a matrix.
• Timestep Embedding: A clock-like mechanism that informs the model which training step it is currently on and how much sharper it needs to generate the image.

Interestingly, another user experienced the exact same issue, not with FLUX, but with SDXL. This proves that the bug is not isolated to a specific model, but is rather an error that triggers inside embeddings.py when complex options are enabled simultaneously across multiple models.

This user noted: “The error only appears after 836 training steps.” This indicates that their training was running perfectly fine initially, but crashed the moment it attempted to pass through the embedding layer (linear_1).

Actual Code Analysis and Debugging: Implementing the PR Fix

Inside embeddings.py, there are two specific blocks (classes) that utilize linear_1:

1. **TimestepEmbedding:** The block that processes timestep information (the clock tracking image generation stages). This is where the second user's SDXL error (2x7424 and 2816x1280) exploded. The error triggered because the batch size of the data mismatched when the model attempted to generate validation images during training.
2. **PixArtAlphaTextProjection:** The block that processes prompt text instructions. This is where the first user's FLUX model error (2x1536 and 768x3072) occurred. The FLUX model internally shares and uses the PixArt structure for text projection.

When a user enables the --with_prior_preservation option, the system handles two types of datasets simultaneously, doubling the data volume. The moment this doubled data enters either TimestepEmbedding or PixArtAlphaTextProjection, the matrix multiplication fails and crashes.

To fix this, I added code in both blocks to compare the columns of linear_1 with the rows of the incoming data before passing it into linear_1. If they mismatch, the code checks the shape of the data. If the data has doubled for any reason, the code splits it in half and passes it through.

I created a new helper function called _unstack_doubled_features:

Python

def _unstack_doubled_features(tensor: torch.Tensor, expected_features: int) -> torch.Tensor:
    if tensor.shape[-1] == expected_features * 2:
        first, second = tensor.chunk(2, dim=-1)
        return torch.cat([first, second], dim=0)
    return tensor

When a user enables the --with_prior_preservation option and a horizontally stretched tensor like [2, 1536] comes in, this function splits the width in half (dim=-1). It creates two separate [2, 768] tensors and then stacks them vertically using torch.cat.

Testing the Implementation

I ran tests using pytest to verify the logic:

• **torch.chunk(sample, 2, dim=-1):** sample is the original data, 2 is the number of pieces to split it into, and dim is the dimension direction. Since PyTorch shapes follow [rows, columns], dim=-1 points to the last dimension (columns), meaning it slices the tensor horizontally (splitting it left and right).
• “Pass both chunks through the sequential layers (linear_1 -> act -> linear_2 -> post_act)": act stands for the activation function. Because a model cannot learn complex patterns through simple matrix multiplication and addition alone, an activation function applies non-linear mathematical operations to allow complex pattern learning. Functions like ReLU, SiLU, and GELU are used here.
• **cat** stands for concatenate.

Python

def forward(self, caption):
    # If prior preservation doubles the feature dimension horizontally
    if caption.shape[-1] == self.linear_1.in_features * 2:

• **forward:** The forward pass. Data moves forward sequentially through mathematical operations. This is a standard PyTorch method.
• **self:** Refers to the class instance itself. Used to access the layers or configuration settings belonging to this specific block.
• **caption:** The incoming data. The prompt is converted into numbers, forming a caption matrix that enters this block.
• **in_features:** The standard, originally designed horizontal dimension size.

Python

chunk1, chunk2 = caption.chunk(2, dim=-1)
        caption = torch.cat([chunk1, chunk2], dim=0)

The incoming caption data is sliced horizontally into two pieces, named chunk1 and chunk2. Then, using the PyTorch cat function, these two chunks are stacked vertically (along dim=0) and saved back into caption.

Python

hidden_states = self.linear_1(caption)
    hidden_states = self.act_1(hidden_states)
    hidden_states = self.linear_2(hidden_states)

The vertically stacked caption is passed through the model's linear_1 layer, and the output is stored in hidden_states. Next, because simply adding and multiplying matrices limits the model's capability, the result from linear_1 is fed into the model's act_1 activation function to enable more complex representations. Finally, it is passed through the linear_2 layer.

The Third PR Attempt: Git Workflows, Tensor Dimensions, and FLUX Architecture

Git Commands for the PR Process

1. Before starting any new work, I need to update the diffusers code on my local machine with the latest code from Hugging Face. Since I might be stuck in a previous branch, I return to the main branch: git checkout main Then, fetch the latest changes from the official Hugging Face repository: git fetch upstream Synchronize my local main branch with the upstream main state: git reset --hard upstream/main
2. Create a isolated branch dedicated only to the new task: git checkout -b my-desired-branch-name
3. After finishing the code modifications, push the changes to my remote repository: git add . / git commit -m "my commit message" / git push origin my-desired-branch-name
4. Submitting the PR on GitHub: Once pushed, if I navigate to my forked diffusers repository on the GitHub website, a green "Compare & pull request" button appears. Click it, write the title and description, and submit.

Another Issue: Another Tensor Shape and Model Dimension Conflict

Another user reported that a specific data tensor passing through the model is triggering a dimensional mismatch.

The linear layers and attention mechanisms inside the model expect fixed input dimensions. However, as the data passes through multiple layers, its size shifts — some parts becoming 4608 dimensions while others become 5120 dimensions — causing the operations to fail.

=> I will need to track down tensor a and tensor b in the source code and align their shapes.

Tensor dimensions start from index 0. For example, in a 3D tensor structured as [Batch, Sequence_Length, Hidden_Dimension], the indices correspond to dim 0, dim 1, and dim 2 respectively (dim 0 is batch size, dim 1 is token sequence length, and dim 2 is feature dimension).

• Singleton Dimension: A dimension with a size of 1 (e.g., [1, 12, 5120]). In PyTorch, broadcasting technology automatically expands this size during calculations.
• Non-singleton Dimension: A fixed dimension with a size greater than 1.

Understanding DreamBooth and Prior Preservation

• DreamBooth: A fine-tuning technique that takes a few photos of a specific subject and trains the model to generate that subject in diverse angles, poses, or backgrounds.
• Prior Preservation: Protects against overfitting and language drift to preserve prior knowledge.
• Language Drift: A phenomenon where an AI, as it trains further, gradually deviates from human language patterns and begins outputting incoherent text. During training, prior/class images are bundled alongside the instance data.

The user is reporting that the error occurs specifically when using the FLUX model with the --with_prior_preservation option enabled.

FLUX is an exceptionally tricky model when it comes to handling tensor dimensions. It utilizes a Dual Text Encoder system, running CLIP L (fixed output dimension) and T5 XXL (fixed output dimension) simultaneously.

• CLIP L: Aligns images and their corresponding words into the same vector space, making it highly effective at understanding image-text pairs.
• T5 XXL: Comprehends text with extreme granularity, grasping complex context and syntax.

Furthermore, FLUX employs a Double Stream Block architecture. It features a unique attention mechanism that concatenates (cat) the text embedding dimension and the image vector dimension together during processing.

Deep Dive into RoPE Mismatch: Vector Dimensions, Multi-GPU Dynamics, and the Real Culprit

The Core Mechanism and Dimensional Conflict

During standard image training, text embeddings and image tensors enter the pipeline perfectly aligned. However, enabling the --with_prior_preservation option bundles the instance data and class data into a single batch before feeding them into the encoder. In this process, a failure occurs because the tensors fail to match the final projection dimension expected by the model.

Reproduction Arguments & System Configuration

To completely reproduce this bug, the following settings and environment variables are used:

• Prompts: --class_prompt="A photo" vs --instance_prompt="A sks photo". The regularization prompt consists of two words, while the training prompt consists of three. This forces the text encoder to output different token sequence lengths, highly likely triggering a dimension conflict when concatenating the two tensors.
• Batch Size: --train_batch_size=1. Although explicitly set to 1, the internal batch size practically becomes 2 because of the simultaneous processing of instance and class data.
• Precision: --mixed_precision="bf16". The numerical computation runs on the BFloat16 data type.

The user pointed out a specific tensor shape discrepancy between 4608 and 5120:

• 4608: This value is roughly aligned with the combined dimensions of FLUX’s text encoders: the T5 embedding (4096 dimensions) plus the CLIP embedding (768 dimensions).
• 5120: This is the exact input dimension expected by FLUX’s internal Double Stream Block or specific linear layers.

Traceback & Execution Log Analysis

The failure trail is documented as follows:

1. The execution begins at line 1720 of train_dreambooth_lora_flux.py. Data is injected into the FLUX transformer model via model_pred = transformer(...).
2. The execution routes through transformer_flux.py into attention_processor.py, where FLUX's attention mechanism computes operations.
3. The script explodes at line 1204 of embeddings.py:
4. out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
5. During the application of RoPE (Rotary Position Embedding) via apply_rotary_emb, the shape mismatch between tensor x and the cos/sin tensors causes the crash.

The FLUX architecture calculates the RoPE position embeddings by combining the image tensor and the text tensor. Consequently, the second axis (dimension 1, the sequence length axis) becomes the sum of the image token count and the text token count. While the error log flags an issue at the third axis (dimension 2), the actual shape mismatch is pushed down during PyTorch’s internal flattening and broadcasting operations, representing a classic token length mismatch.

When --with_prior_preservation is active, both instance and class images are processed. However, under the user's settings:

• Instance Prompt: "A sks photo" (Token count + Padding)
• Class Prompt: "A photo" (Token count + Padding)

These two distinct prompts should pass through the text encoder and align cleanly into a single batch (via padding and concatenation). Instead, the text embedding logic is missing or incompletely implemented. This leaves the instance data tensor configuration at 4608 and the class data configuration at 5120. Because their lengths differ, the system crashes the exact moment it attempts to apply the RoPE position embeddings (cos, sin).

System Info & Multi-GPU Infrastructure

• Diffusers Version: 0.33.0.dev0. This indicates a developer version built directly on the main branch and latest codebase.
• PyTorch Version: 2.5.1+cu124. Tensor dimension checking was strictly reinforced for apply_rotary_emb operations in this version, blocking broadcasting and triggering the RuntimeError.
• Accelerator: NVIDIA A800-SXM4-80GB. The user is running training across a distributed multi-GPU setup with 8 high-performance GPUs. When prior preservation is enabled in a multi-GPU environment, the accelerate library splits and distributes the batch across processes via split_between_processes. During this step, either the instance and class data batches were unevenly split, or the text encoder outputs heavily stacked on one specific GPU, fracturing the tensor dimensions.

Finding the Real Culprit

Python

if args.with_prior_preservation:
    prompt_embeds = torch.cat([prompt_embeds, class_prompt_hidden_states], dim=0)
    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, class_pooled_prompt_embeds], dim=0)
    text_ids = torch.cat([text_ids, class_text_ids], dim=0) # <--- The Culprit

FLUX processes text and images together by combining them into a single sequence, relying on text_ids for text positioning and image_ids for image positioning.

The code snippet above concatenates the two tensors vertically along dim=0 (the batch axis). However, because text_ids is stored as a 2D tensor structured as [length, 3], performing torch.cat along dim=0 inadvertently expands the sequence length itself—stretching the shape from (512, 3) to (1024, 3). This causes a structural distortion.

Instead of scaling the batch dimension, the code artificially inflates the sequence size. When the model tries to map this to internal RoPE calculations, it cannot reconcile the distorted shape with the expected 4608 vs 5120 position embeddings, leading to a crash.

Evaluation of Alternative Workarounds

Another contributor attempted a fix by re-encoding the prompt at each step:

Python

else:
    elems_to_repeat = len(prompts)
    # ... truncated ...
    # Re-encodes the prompt by repeating it to match the batch size at every training step

This workaround checks the total count of input prompts when --with_prior_preservation is on, generating new text embeddings and matching text_ids to force manual synchronization.

• Limitation: This is a suboptimal patch. Re-running the text encoder redundantly at every single step significantly compromises training speeds.

The Elegant Solution

The correct shape FLUX expects for text_ids is [batch size, sequence length, 3], or it must maintain an independent dimension structure for each batch. The reason torch.cat mistakenly altered the sequence axis instead of the batch axis is that the tensor was 2D [length, 3], making dim=0 point to length. To fix this cleanly, one would typically expand the dimensions using None or change the structure using torch.stack.

However, in FLUX, text_ids serves as a tensor containing text token positional index data. If all data within the incoming batch shares an identical structure and length, FLUX allows the batches to share the exact same text_ids positional data, automatically handling broadcasting internally.

The buggy code took the identical class_text_ids and instance_text_ids and forcefully concatenated them into a doubled tensor just because the batch size doubled. The model architecture expects a single-fold text_ids size, so throwing a doubled size into the RoPE layers breaks the dimensions.

The fix is simple: just delete the concatenation line.

Python

text_ids = torch.cat([text_ids, class_text_ids], dim=0) # Delete this line

Because class_text_ids shares the exact same shape and dimensions as instance_text_ids, there is absolutely no need to distort the shape using torch.cat. Leaving instance_text_ids as it is allows the core engine to broadcast and apply the dimensions correctly across the batch.

Three Code Paths in the FLUX Script

The FLUX script routes execution through three pathways depending on prompt configurations and encoder training settings:

• custom_instance_prompts=True: When custom prompts are used, the correct text_ids compliant with the model architecture is handled naturally.
• train_text_encoder=True: During text encoder training paths, after tokens are merged, the script cleanly regenerates the text_ids shape inside the encode_prompt function using torch.zeros(prompt_embeds.shape[1], 3).
• Fixed Default Prompts with Prior Preservation: When neither of the above conditions is met and only the Prior Preservation option is enabled, the bug strikes at line 1536. This is where text_ids = torch.cat([text_ids, class_text_ids], dim=0) forces a broken merge between class_text_ids and instance_text_ids.