File size: 2,756 Bytes
c750327 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | # ==============================================================================
# BlockDiffuse Improved GPU Profile (Transfer Learning, CE, NN Loss, Curriculum)
# ==============================================================================
# Base LLM Metadata
base_llm:
model_name_or_path: "Qwen/Qwen2.5-0.5B-Instruct"
mid_layer_idx: 12
d_model: 896
num_layers: 24
num_heads: 14
vocab_size: 151936
max_prompt_len: 512
max_target_len: 100
# High-Capacity DiT Architecture with Transfer Learning & Deep Projection Head
dit:
d_model: 896
num_layers: 10 # Deeper 8-layer DiT model
num_heads: 14 # 14 attention heads (head_dim = 64)
mlp_ratio: 8.0
block_length: 100 # Maximum 100-token parallel generation
max_seq_len: 1024
dropout: 0.1
use_rope: true
use_block_causal: false
gradient_checkpointing: true # Dramatically cuts activation memory (~2GB VRAM)
adaln_zero: true
time_embedding_dim: 256
apply_rmsnorm_head: true
use_projection_head: true
projection_head_depth: "deep" # 3-layer deep MLP with 4x hidden expansion
init_from_base: true # Layer initialization from Qwen mid-layers (6-11)
init_num_layers: 6
init_start_layer: 6
init_base_model_name: "Qwen/Qwen2.5-0.5B-Instruct"
# Training Hyperparameters with Block-Length Curriculum
training:
learning_rate: 7.0e-4
min_lr: 1.0e-6
weight_decay: 0.01
warmup_steps: 200
max_steps: 5000
batch_size: 16 # 4x reduced batch size (effective batch = 8 with grad accum 2)
gradient_accumulation_steps: 1 # Effective global batch size = 32
max_grad_norm: 1.0
precision: "bfloat16" # High performance native bfloat16
save_interval: 1000
eval_interval: 200
log_interval: 10
output_dir: "./checkpoints_improved"
# Gradual Block-Length Curriculum
block_length_curriculum: true
curriculum_start_steps: 0
curriculum_increase_interval: 2000
curriculum_factor: 2
curriculum_min_block: 20
# Multi-Objective Loss with Discrete Token CE & Contrastive NN Supervision
loss_weights:
lambda_fm: 1.0 # Rectified flow matching velocity MSE (L_FM)
lambda_disp: 0.1 # Dispersive variance regularizer (L_Disp)
lambda_kl: 0.1 # Teacher Logit Distillation Loss (L_KL)
use_kldiv_loss: true
kl_temperature: 2.0
lambda_ce: 1.0 # Discrete Cross-Entropy Loss on Tokens (L_CE)
use_ce_loss: true
lambda_nn: 0.1 # Contrastive InfoNCE Nearest-Neighbor Loss (L_NN)
use_nn_loss: true
nn_temperature: 0.1
flow_matching:
sigma_min: 1.0e-5
time_sampling: "uniform"
|