LLaDA2.0 Uni FP8

--- license: apache-2.0 language: - en tags: - multimodal - image-generation - image-understanding - image-editing - diffusion - moe - text-to-image - fp8 - quantized library_name: transformers pipeline_tag: any-to-any base_model: inclusionAI/LLaDA2.0-Uni --- <p align="center"> <img src="./assets/llada_logo.png" width="20%"/> </p> <div align="center"> <h1> LLaDA2.0-Uni: Unifying Multimodal Understanding and Generation with Diffusion Large Language Model </h1> <h3> FP8 Quantized Version of LLaDA2.0-Uni </h3> [[📑 Technical Report]()] &emsp; [[🌐 Github](https://github.com/inclusionAI/LLaDA2.0-Uni)] <b>AGI Research Center, Inclusion AI</b> </div> ## Overview This is the **FP8 quantized** version of [LLaDA2.0-Uni](https://huggingface.co/inclusionAI/LLaDA2.0-Uni), featuring block-wise FP8 quantization of MoE expert weights. This reduces GPU memory usage by ~48% for model loading while preserving output quality. ### Quantization Details - **Method**: Block-wise FP8 (float8_e4m3fn) with per-block scale factors - **Block size**: 128×128 - **Quantized layers**: MoE routed expert weights (gate_proj, up_proj, down_proj) - **Kept in BF16**: Embeddings, lm_head, attention projections, shared experts, layer norms, routing gates ### Memory Comparison | Variant | Model Loading | T2I Peak | Understanding Peak | Edit Peak | |---------|--------------|----------|-------------------|-----------| | BF16 | 62.9 GB | 35.3 GB | 33.2 GB | 41.7 GB | | FP8 | 32.5 GB | 35.3 GB | 33.3 GB | 41.8 GB | > **Note**: FP8 halves the static model weight memory (~30 GB saved at load time). Peak inference memory is similar because activations dominate during generation. ## Quick Start ```python from transformers import AutoModelForCausalLM, AutoTokenizer model_path = "inclusionAI/LLaDA2.0-Uni-FP8" tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( model_path, device_map="cuda", trust_remote_code=True ).eval() model.tokenizer = tokenizer # Text-to-Image Generation result = model.generate_image( "A cat sitting on a windowsill at sunset", image_h=1024, image_w=1024, steps=16, cfg_scale=4.0, ) # Decode VQ tokens to image from decoder import decode_vq_tokens image = decode_vq_tokens( result["token_ids"], result["h"], result["w"], model_path, "cuda", num_steps=8, decode_mode="decoder-turbo", ) image....