[RFC]: Fix + Standardize DiT Config/Init Patterns

[alex-jw-brooks]

Motivation.

This will fix existing bad behaviors in DiT initialization around default values from Diffusers and hopefully add a sustainable pattern for well-typed DiT configs - this will also handle merging parsed config values with default values in a way that keeps the model code as clear as possible.

Proposed Change.

This is mostly an internal change for preventing bugs in config parsing / DiT initialization and making the code more clear, so there will be no changes to accuracy, performance, or APIs.

Summary of Current Problem

There are currently some issues with how configs are being handled, which partially stems from the way in which Diffusers manages its configs compared to Transformers (i.e., through the use of mapping to classes via mixins rather than having easily inspectable PretrainedConfig classes per architecture).

In the current pipeline for Diffusion models, most components are created .from_pretrained, but the Transformer component is not initialized in a way that is consistent. There are generally 3 ways that we initialize it:

1. By passing the omni diffusion config and pulling things off thge DiT config from in the class
2. By exploding the DiT into kwargs derived from the DiT config
3. By relying on defaults in the __init__ of the transformer class

The last of which is particularly problematic, because it causes classes to be initialized incorrectly. The defaults are copied from the corresponding Diffusers class, but since these tend to be overridden from the subconfigs they're loaded, they actually may not be correct.

Example - LongCat

The current LongCat config management illustrates this issue well;

• In the current Image pipeline, we pass the Diffusion config here

self.transformer = LongCatImageTransformer2DModel(od_config=od_config)

While the intializer takes many of these values here:

    def __init__(
        self,
        od_config: OmniDiffusionConfig,
        patch_size: int = 1,
        in_channels: int = 64,
        num_layers: int = 19,
        num_single_layers: int = 38,
        attention_head_dim: int = 128,
        num_attention_heads: int = 24,
        joint_attention_dim: int = 3584,
        pooled_projection_dim: int = 3584,
        axes_dims_rope: list[int] = [16, 56, 56],
    ):

whose kwargs are used to set the attributes on the class, not the passsed transformer config. These values are understandably taken from the default of the corresponding Diffusers class here.

However, AFAIK they are not the right values for any LongCat model; I think they are copied from Flux as the reference architecture. E.g,. the Transformers config for LongCatImage.

{
    "_class_name": "LongCatImageTransformer2DModel",
    "_diffusers_version": "0.30.0.dev0",
    "attention_head_dim": 128,
    "guidance_embeds": false,
    "in_channels": 64,
    "joint_attention_dim": 3584,
    "num_attention_heads": 24,
    "num_layers": 10,
    "num_single_layers": 20,
    "patch_size": 1,
    "pooled_projection_dim": 3584
}

As such, the pipeline ends up incorrectly configured, and in this case much heavier than it should be, because it initializes extra transformer blocks that it doesn't have weights for. I noticed this while working on TeaCache support for LongCat since I was seeing nans and unexpected OOMS while calibrating coefficients through the transformers block. This is the approach currently used for diffusion pipelines.

Proposed Solution

Add a configs dir vllm_omni.diffusion.config which implements a well-defined PretrainedConfig subclass per model, which handles its default kwargs, and can be used for type hints, similar to the custom transformers configs in vLLM here.

More concretely, a small base class could be added for converting from the generic transformers config wrapper:

class BaseDiTConfig(ABC, PretrainedConfig):
    _class_name: str | None = None

    @classmethod
    def from_tf_config(cls, cfg: TransformerConfig):
        """For now - converts the generic wrapper around Transformers Configs that is currently used for DiT
        configs to an instance of this subclass.
        """
        model_dict = cfg.to_dict()
        <validate against the _class_name as needed>
        return cls.from_dict(model_dict)

With something like this per DiT class

from vllm_omni.diffusion.config.base import BaseDiTConfig

class ZImageTransformer2DModelConfig(BaseDiTConfig):
    # Expected _class_name in Diffusers
    _class_name = "ZImageTransformer2DModel"

    def __init__(
        self,
        all_patch_size=(2,),
        all_f_patch_size=(1,),
        ...
        axes_dims=[32, 48, 48],
        axes_lens=[1024, 512, 512],
        **kwargs,
    ):
        super().__init__()
        self.all_patch_size = all_patch_size
        self.all_f_patch_size = all_f_patch_size
        ...
        self.axes_dims = axes_dims
        self.axes_lens = axes_lens

And in the corresponding pipeline:

        hf_config = ZImageTransformer2DModelConfig.from_tf_config(od_config.tf_model_config)
        self.transformer = ZImageTransformer2DModel(
            hf_config=hf_config,
            od_config=od_config, # Should also be passed for things like quant config etc
        )

I think that passing the pretrained config subclass directly, as opposed to building it from the omni diffusion config in the DiT is probably a bit cleaner when considering some of the model architectures, e.g., wan2, which may need to build multiple transformers in its pipeline.

        self.transformer = get_dit_model(
            od_config=od_config,
            subfolder="transformer",
            local_files_only=local_files_only,
        ) if load_transformer else None

        self.transformer_2 = get_dit_model(
            od_config=od_config,
            subfolder="transformer_2",
            local_files_only=local_files_only,
        ) if load_transformer_2 else None

I think this would also be nice for potential follow-ups, e.g,. can push the config conversions a bit earlier on in the process to avoid the current generic transformers wrapper and make the pipeline download calls more well-patterned in the future

Feedback Period.

One week - will also open a draft for discussion in the coming couple of days!

CC List.

@DarkLight1337 @Isotr0py @ywang96 @sfeng33 @NickLucche @hsliuustc0106 @Gaohan123 @tzhouam @ZJY0516 in case anyone has thoughts

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[DarkLight1337]

Let me ask @hmellor about this, see if it's possible to fix this from HF side.

[alex-jw-brooks]

Cool thanks @DarkLight1337 - I spent some time on this today and opened a draft PR for a few models here. CC'ed you both on it in case it's easier to discuss concretely

[hmellor]

This is really valuable feedback and definitely something that seems would be better fixed on the Diffusers side.

I've forwarded this RFC to the Diffusers team, I'll let you know their response if they don't do it directly here.

[sayakpaul]

Hey folks,

Sayak here from the Diffusers team 👋

Currently, on diffusers, we can do:

from diffusers import AutoModel

ckpt_id = "black-forest-labs/FLUX.1-dev"
config = AutoModel.load_config(ckpt_id, subfolder="transformer")
model = AutoModel.from_config(config)

Of course, one can change this config object to pass custom values.

Does this solve this inconsistency in the caller sites?

We can work on doing better defaults, but note that most of these are taken directly from the original checkpoints. Some inconsistencies can arise due to the fact that Flux2Transformer2DModel can be used for both Flux2 and Flux2 Klein (although the difference here is just in the number of layers).

[alex-jw-brooks]

Hey @sayakpaul, thanks for the quick response! I think there are two issues this RFC intends to address:

1. Standardizing the pattern for how DiTs in vLLM Omni are initialized and removing existing bugs that exist because it's not well-patterned yet
2. Making model development for DiTs in vLLM Omn well-patterned & cleaner

The second one is the main intention with adding well-defined configs classes, because underneath, we are currently doing pretty much the same thing, i.e.,

            ...
            tf_config_dict = get_hf_file_to_dict("transformer/config.json", od_config.model) # `json.load`s the downloaded file
            od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict) # takes the dict and puts it in a generic wrapper

The main advantages for having the concrete types with defaults set is the following:

1. Clarity & type hints for contributors writing the DiTs
2. It decouples the vLLM Omni DiT implementations from the Diffusers implementations, which will make managing versions easier
   • E.g., if the Diffusers DiT implementation modified config key(s), it could be cleanly handled to remap correctly to Diffusers. However, since we are just loading a dict in this case and initializing our own DiT with it, we would need to handle this correctly on our side atm, which may not be obvious without clear typing
3. All values are well-defined; currently we have some things like this in DiTs

        self.dual_attention_layers = (
            model_config.dual_attention_layers if hasattr(model_config, "dual_attention_layers") else ()
        )
        self.qk_norm = model_config.qk_norm if hasattr(model_config, "qk_norm") else ""

This is not ideal, because in the refactor example above, we would not be able to easily catch something like dual_attention_layers is not in the config, so we should use () vs dual_attention_layers was renamed to something else, but since it isn't there, we use the default value of () anyway 🙂 This is another reason why something like this is IMO much more maintainable, because we can more easily add validation to check for potential key misalignment and not have to rely on hasattr/getattr for things that may be excluded in the raw config file

[sayakpaul]

Thanks for the continued and elaborate discussions! PFB my thoughts.

The main advantages for having the concrete types with defaults set is the following:

We can work with the vllm-omni team to improve typing so that it's a better developer experience. Would you be willing to pair with me on this? For defaults, I believe we already have sensible defaults on most of our implementations (at least the modern ones from the last ~12 months).

One exception is the case where a modeling class like Flux2Transformer2DModel is used for both Flux2 and Flux2 Klein. This is how the original codebase also implements it and IMO, it makes sense because they both differ just in terms of the number of layers.

E.g., if the Diffusers DiT implementation modified config key(s), it could be cleanly handled to remap correctly to Diffusers. However, since we are just loading a dict in this case and initializing our own DiT with it, we would need to handle this correctly on our side atm, which may not be obvious without clear typing

I am not too sure about it. It would be helpful if you could provide a representative example of this. If better typing on diffusers would help alleviate this issue, we're of course, up for it.

All values are well-defined; currently we have some things like this in DiTs

I agree this is not ideal and we're happy to standardize on this on diffusers. I believe this has largely been the case for the models added in the last ~12 months in diffusers. If there's bad residue, we should fix that. Could you share some examples where this is happening?

In general, init arguments are rarely renamed in diffusers unless we have an extremely strong reason to do so. We're happy to keep you folks in the loop if and when those cases arise in the future. This way, we can brainstorm the solutions together. WDYT?

[hsliuustc0106]

We can work with the vllm-omni team to improve typing so that it's a better developer experience. Would you be willing to pair with me on this? For defaults, I believe we already have sensible defaults on most of our implementations (at least the modern ones from the last ~12 months).

cc @ZJY0516 @wtomin

[alex-jw-brooks]

For sure @sayakpaul, thanks! LongCat Image is a good example for both unhappy defaults and not setting all of the kwargs in diffuseres; it has wrong defaults for num_single_layers / num_layers , and the config doesn't set axes_dim_ropes to pass here.

The layers thing is an example of a currently existing copy paste from our side that will fixed by this, but axes_dim_ropes isn't in the resulting dict when you load the config for example.

ckpt_id = "meituan-longcat/LongCat-Image"
config = AutoModel.load_config(ckpt_id, subfolder="transformer")

{'_class_name': 'LongCatImageTransformer2DModel', '_diffusers_version': '0.30.0.dev0', 'attention_head_dim': 128, 'guidance_embeds': False, 'in_channels': 64, 'joint_attention_dim': 3584, 'num_attention_heads': 24, 'num_layers': 10, 'num_single_layers': 20, 'patch_size': 1, 'pooled_projection_dim': 3584}

Maybe an easy place to start on the Diffusers side that we could potentially iterate on would be a utility that

1. Gets the the config as above ^
2. Uses the retrieved config["_class_name"] to determine the DiT class being initialized
3. inspects the retrieved class for default kwargs
4. returns the {**config, **model_defaults} to apply the overrides and missing keys

This does unfortunately keep things as a raw dict, but at least would ensure incorrect defaults are overridden in the same way, and that all expected kwargs to init are in the dict if we have want to know what the DiT actually would be using if the same config were used in Diffusers. I am also happy to open a PR in Diffusers to discuss, if it would help!

[sayakpaul]

This is all very reasonable and detailed feedback. I want to start by acknowledging that. Thank you! Please keep doing this :)

Sometimes, during the integration process, the authors themselves lead them and we might have to prioritize speed for supporting the release (which I believe was the case for LongCat). This likely explains the reason why you saw it for LongCat but not for other modern models. But this doesn't excuse the friction that gets introduced in the developer experience (+ the inconsistency).

LongCat Image is a good example for both unhappy defaults and not setting all of the kwargs in diffuseres; it has wrong defaults for num_single_layers / num_layers , and the config doesn't set axes_dim_ropes to pass here.

We can definitely open a PR to fix that config and work with the author to get that merged.

This does unfortunately keep things as a raw dict, but at least would ensure incorrect defaults are overridden in the same way, and that all expected kwargs to init are in the dict if we have want to know what the DiT actually would be using if the same config were used in Diffusers. I am also happy to open a PR in Diffusers to discuss, if it would help!

Of course, we will be quick to review and work with you to get that prioritized.

[DN6]

Hi @alex-jw-brooks. I think there are a couple of options to get the required config.

If you load the model from_config on meta device, you could access the config from the model instance without the overhead of creating weights.

import torch
from diffusers import AutoModel

repo = "meituan-longcat/LongCat-Image"
subfolder = "transformer"

config = AutoModel.load_config(repo, subfolder=subfolder)

with torch.device("meta"):
    model = AutoModel.from_config(config)
print(f"model.config:")
for k, v in dict(model.config).items():
    if not k.startswith("_"):
        print(f"  {k}: {v}")

Would print

model.config:
  patch_size: 1
  in_channels: 64
  num_layers: 10
  num_single_layers: 20
  attention_head_dim: 128
  num_attention_heads: 24
  joint_attention_dim: 3584
  pooled_projection_dim: 3584
  axes_dims_rope: [16, 56, 56]
  guidance_embeds: False

The keys and value types would correspond with exactly what the init expects.

We could explore having a utility method that returns a typed dataclass config for the model. e.g.

from diffusers import AutoModel

config_dict = AutoModel.load_config(
    "meituan-longcat/LongCat-Image",
    subfolder="transformer",
)
# import DiT based on _class_name
typed_config = LongCatImageTransformer2DModel._create_dataclass_from_config(config_dict)
for f in dataclasses.fields(typed_config):
     print(f"{f.name}: {f.type}")

But I think inspecting the key names and types in the FrozenDict that we currently return with the config property would be just as effective as a dataclass since all the init args have to be JSON serializable anyway.