Text Generation
Diffusers
English
diffusion
flow-matching
rectified-flow
reasoning
qwen2.5
block-diffusion
non-autoregressive
Instructions to use tahamajs/BlockDiffuse with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Diffusers
How to use tahamajs/BlockDiffuse with Diffusers:
pip install -U diffusers transformers accelerate
import torch from diffusers import DiffusionPipeline # switch to "mps" for apple devices pipe = DiffusionPipeline.from_pretrained("tahamajs/BlockDiffuse", dtype=torch.bfloat16, device_map="cuda") prompt = "Hi, what can you help me with?" image = pipe(prompt).images[0] - Notebooks
- Google Colab
- Kaggle
Upload config.yaml with huggingface_hub
Browse files- config.yaml +76 -0
config.yaml
ADDED
|
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ==============================================================================
|
| 2 |
+
# BlockDiffuse Improved GPU Profile (Transfer Learning, CE, NN Loss, Curriculum)
|
| 3 |
+
# ==============================================================================
|
| 4 |
+
|
| 5 |
+
# Base LLM Metadata
|
| 6 |
+
base_llm:
|
| 7 |
+
model_name_or_path: "Qwen/Qwen2.5-0.5B-Instruct"
|
| 8 |
+
mid_layer_idx: 12
|
| 9 |
+
d_model: 896
|
| 10 |
+
num_layers: 24
|
| 11 |
+
num_heads: 14
|
| 12 |
+
vocab_size: 151936
|
| 13 |
+
max_prompt_len: 512
|
| 14 |
+
max_target_len: 100
|
| 15 |
+
|
| 16 |
+
# High-Capacity DiT Architecture with Transfer Learning & Deep Projection Head
|
| 17 |
+
dit:
|
| 18 |
+
d_model: 896
|
| 19 |
+
num_layers: 10 # Deeper 8-layer DiT model
|
| 20 |
+
num_heads: 14 # 14 attention heads (head_dim = 64)
|
| 21 |
+
mlp_ratio: 8.0
|
| 22 |
+
block_length: 100 # Maximum 100-token parallel generation
|
| 23 |
+
max_seq_len: 1024
|
| 24 |
+
dropout: 0.1
|
| 25 |
+
use_rope: true
|
| 26 |
+
use_block_causal: false
|
| 27 |
+
gradient_checkpointing: true # Dramatically cuts activation memory (~2GB VRAM)
|
| 28 |
+
adaln_zero: true
|
| 29 |
+
time_embedding_dim: 256
|
| 30 |
+
apply_rmsnorm_head: true
|
| 31 |
+
use_projection_head: true
|
| 32 |
+
projection_head_depth: "deep" # 3-layer deep MLP with 4x hidden expansion
|
| 33 |
+
init_from_base: true # Layer initialization from Qwen mid-layers (6-11)
|
| 34 |
+
init_num_layers: 6
|
| 35 |
+
init_start_layer: 6
|
| 36 |
+
init_base_model_name: "Qwen/Qwen2.5-0.5B-Instruct"
|
| 37 |
+
|
| 38 |
+
# Training Hyperparameters with Block-Length Curriculum
|
| 39 |
+
training:
|
| 40 |
+
learning_rate: 7.0e-4
|
| 41 |
+
min_lr: 1.0e-6
|
| 42 |
+
weight_decay: 0.01
|
| 43 |
+
warmup_steps: 200
|
| 44 |
+
max_steps: 5000
|
| 45 |
+
batch_size: 16 # 4x reduced batch size (effective batch = 8 with grad accum 2)
|
| 46 |
+
gradient_accumulation_steps: 1 # Effective global batch size = 32
|
| 47 |
+
max_grad_norm: 1.0
|
| 48 |
+
precision: "bfloat16" # High performance native bfloat16
|
| 49 |
+
save_interval: 1000
|
| 50 |
+
eval_interval: 200
|
| 51 |
+
log_interval: 10
|
| 52 |
+
output_dir: "./checkpoints_improved"
|
| 53 |
+
|
| 54 |
+
# Gradual Block-Length Curriculum
|
| 55 |
+
block_length_curriculum: true
|
| 56 |
+
curriculum_start_steps: 0
|
| 57 |
+
curriculum_increase_interval: 2000
|
| 58 |
+
curriculum_factor: 2
|
| 59 |
+
curriculum_min_block: 20
|
| 60 |
+
|
| 61 |
+
# Multi-Objective Loss with Discrete Token CE & Contrastive NN Supervision
|
| 62 |
+
loss_weights:
|
| 63 |
+
lambda_fm: 1.0 # Rectified flow matching velocity MSE (L_FM)
|
| 64 |
+
lambda_disp: 0.1 # Dispersive variance regularizer (L_Disp)
|
| 65 |
+
lambda_kl: 0.1 # Teacher Logit Distillation Loss (L_KL)
|
| 66 |
+
use_kldiv_loss: true
|
| 67 |
+
kl_temperature: 2.0
|
| 68 |
+
lambda_ce: 1.0 # Discrete Cross-Entropy Loss on Tokens (L_CE)
|
| 69 |
+
use_ce_loss: true
|
| 70 |
+
lambda_nn: 0.1 # Contrastive InfoNCE Nearest-Neighbor Loss (L_NN)
|
| 71 |
+
use_nn_loss: true
|
| 72 |
+
nn_temperature: 0.1
|
| 73 |
+
|
| 74 |
+
flow_matching:
|
| 75 |
+
sigma_min: 1.0e-5
|
| 76 |
+
time_sampling: "uniform"
|