| """Configuration objects for the Cola2 Full-DiT model.""" |
|
|
| from __future__ import annotations |
|
|
| from typing import Any, Optional |
|
|
| from transformers import GenerationConfig, PretrainedConfig |
|
|
|
|
| _AUTO_MAP = { |
| "AutoConfig": "configuration_cola2.Cola2Config", |
| "AutoModel": "modeling_cola2.Cola2ForConditionalGeneration", |
| "AutoModelForCausalLM": "modeling_cola2.Cola2ForConditionalGeneration", |
| } |
|
|
|
|
| class Cola2Config(PretrainedConfig): |
| """Architecture configuration for span-causal Cola2. |
| |
| Defaults describe the released 4B DiT + Qwen3-4B-shaped decoder. Tests and |
| CPU examples should use :meth:`tiny`, which executes the identical graph at |
| a small width and depth. |
| """ |
|
|
| model_type = "cola2" |
| keys_to_ignore_at_inference = ["past_key_values"] |
|
|
| def __init__( |
| self, |
| vocab_size: int = 151_936, |
| encoder_embedding_size: int = 2_560, |
| encoder_hidden_size: int = 1_536, |
| encoder_intermediate_size: int = 6_144, |
| encoder_num_attention_heads: int = 12, |
| encoder_num_hidden_layers: int = 0, |
| pool_num_hidden_layers: int = 2, |
| latent_encoder_num_hidden_layers: int = 2, |
| latent_decoder_num_hidden_layers: int = 0, |
| latent_dim: int = 64, |
| use_variational_encoder: bool = True, |
| boundary_score_alpha: float = 10.0, |
| boundary_score_threshold: float = 0.5, |
| target_compression_rate: float = 0.5, |
| min_span_length: int = 1, |
| max_span_length: int = 8, |
| latent_block_causal: bool = False, |
| latent_block_size: int = 8, |
| dit_hidden_size: int = 2_560, |
| dit_intermediate_size: int = 10_240, |
| dit_num_hidden_layers: int = 34, |
| dit_num_attention_heads: int = 20, |
| dit_head_dim: int = 128, |
| dit_rope_dim: int = 96, |
| dit_block_size: int = 8, |
| decoder_hidden_size: int = 2_560, |
| decoder_intermediate_size: int = 9_728, |
| decoder_num_hidden_layers: int = 36, |
| decoder_num_attention_heads: int = 32, |
| decoder_num_key_value_heads: int = 8, |
| decoder_head_dim: int = 128, |
| condition_num_attention_heads: int = 20, |
| decoder_condition_num_hidden_layers: int = 2, |
| decoder_lm_head_bias: bool = True, |
| duration_hidden_size: Optional[int] = None, |
| local_position_dim: int = 128, |
| hidden_act: str = "silu", |
| hidden_dropout: float = 0.0, |
| attention_dropout: float = 0.0, |
| layer_norm_eps: float = 1e-6, |
| dit_layer_norm_eps: float = 1e-5, |
| encoder_rope_theta: float = 500_000.0, |
| condition_rope_theta: float = 500_000.0, |
| decoder_rope_theta: float = 1_000_000.0, |
| dit_rope_theta: float = 10_000.0, |
| rope_theta: float = 1_000_000.0, |
| initializer_range: float = 0.02, |
| max_position_embeddings: int = 32_768, |
| scaling_factor: float = 1.0, |
| shifting_factor: float = 0.0, |
| diffusion_time: float = 1_000.0, |
| pad_token_id: int = 151_936, |
| mask_token_id: int = 151_936, |
| bos_token_id: Optional[int] = 151_643, |
| eos_token_id: Optional[int] = 151_643, |
| tie_word_embeddings: bool = False, |
| auto_map: Optional[dict[str, str]] = None, |
| architectures: Optional[list[str]] = None, |
| **kwargs: Any, |
| ) -> None: |
| self.vocab_size = int(vocab_size) |
| self.encoder_embedding_size = int(encoder_embedding_size) |
| self.encoder_hidden_size = int(encoder_hidden_size) |
| self.encoder_intermediate_size = int(encoder_intermediate_size) |
| self.encoder_num_attention_heads = int(encoder_num_attention_heads) |
| self.encoder_num_hidden_layers = int(encoder_num_hidden_layers) |
| self.pool_num_hidden_layers = int(pool_num_hidden_layers) |
| self.latent_encoder_num_hidden_layers = int( |
| latent_encoder_num_hidden_layers |
| ) |
| self.latent_decoder_num_hidden_layers = int( |
| latent_decoder_num_hidden_layers |
| ) |
| self.latent_dim = int(latent_dim) |
| self.use_variational_encoder = bool(use_variational_encoder) |
| self.boundary_score_alpha = float(boundary_score_alpha) |
| self.boundary_score_threshold = float(boundary_score_threshold) |
| self.target_compression_rate = float(target_compression_rate) |
| self.min_span_length = int(min_span_length) |
| self.max_span_length = int(max_span_length) |
| self.latent_block_causal = bool(latent_block_causal) |
| self.latent_block_size = int(latent_block_size) |
|
|
| self.dit_hidden_size = int(dit_hidden_size) |
| self.dit_intermediate_size = int(dit_intermediate_size) |
| self.dit_num_hidden_layers = int(dit_num_hidden_layers) |
| self.dit_num_attention_heads = int(dit_num_attention_heads) |
| self.dit_head_dim = int(dit_head_dim) |
| self.dit_rope_dim = int(dit_rope_dim) |
| self.dit_block_size = int(dit_block_size) |
|
|
| self.decoder_hidden_size = int(decoder_hidden_size) |
| self.decoder_intermediate_size = int(decoder_intermediate_size) |
| self.decoder_num_hidden_layers = int(decoder_num_hidden_layers) |
| self.decoder_num_attention_heads = int(decoder_num_attention_heads) |
| self.decoder_num_key_value_heads = int(decoder_num_key_value_heads) |
| self.decoder_head_dim = int(decoder_head_dim) |
| self.condition_num_attention_heads = int( |
| condition_num_attention_heads |
| ) |
| self.decoder_condition_num_hidden_layers = int( |
| decoder_condition_num_hidden_layers |
| ) |
| self.decoder_lm_head_bias = bool(decoder_lm_head_bias) |
| self.duration_hidden_size = int( |
| duration_hidden_size or max(self.decoder_hidden_size // 2, 64) |
| ) |
| self.local_position_dim = int(local_position_dim) |
|
|
| self.hidden_act = str(hidden_act) |
| self.hidden_dropout = float(hidden_dropout) |
| self.attention_dropout = float(attention_dropout) |
| self.layer_norm_eps = float(layer_norm_eps) |
| self.dit_layer_norm_eps = float(dit_layer_norm_eps) |
| self.encoder_rope_theta = float(encoder_rope_theta) |
| self.condition_rope_theta = float(condition_rope_theta) |
| self.decoder_rope_theta = float(decoder_rope_theta) |
| self.dit_rope_theta = float(dit_rope_theta) |
| self.rope_theta = float(rope_theta) |
| self.initializer_range = float(initializer_range) |
| self.max_position_embeddings = int(max_position_embeddings) |
| self.scaling_factor = float(scaling_factor) |
| self.shifting_factor = float(shifting_factor) |
| self.diffusion_time = float(diffusion_time) |
| self.mask_token_id = int(mask_token_id) |
|
|
| self._validate_architecture() |
| |
| |
| |
| kwargs.pop("is_decoder", None) |
| kwargs.pop("is_encoder_decoder", None) |
|
|
| super().__init__( |
| pad_token_id=int(pad_token_id), |
| mask_token_id=int(mask_token_id), |
| bos_token_id=bos_token_id, |
| eos_token_id=eos_token_id, |
| tie_word_embeddings=bool(tie_word_embeddings), |
| is_decoder=True, |
| is_encoder_decoder=False, |
| auto_map=dict(_AUTO_MAP if auto_map is None else auto_map), |
| architectures=( |
| ["Cola2ForConditionalGeneration"] |
| if architectures is None |
| else architectures |
| ), |
| **kwargs, |
| ) |
|
|
| def _validate_architecture(self) -> None: |
| positive = { |
| "vocab_size": self.vocab_size, |
| "encoder_embedding_size": self.encoder_embedding_size, |
| "encoder_hidden_size": self.encoder_hidden_size, |
| "encoder_intermediate_size": self.encoder_intermediate_size, |
| "encoder_num_attention_heads": self.encoder_num_attention_heads, |
| "latent_dim": self.latent_dim, |
| "latent_block_size": self.latent_block_size, |
| "dit_hidden_size": self.dit_hidden_size, |
| "dit_intermediate_size": self.dit_intermediate_size, |
| "dit_num_attention_heads": self.dit_num_attention_heads, |
| "dit_head_dim": self.dit_head_dim, |
| "dit_block_size": self.dit_block_size, |
| "decoder_hidden_size": self.decoder_hidden_size, |
| "decoder_intermediate_size": self.decoder_intermediate_size, |
| "decoder_num_attention_heads": self.decoder_num_attention_heads, |
| "decoder_num_key_value_heads": self.decoder_num_key_value_heads, |
| "decoder_head_dim": self.decoder_head_dim, |
| "condition_num_attention_heads": self.condition_num_attention_heads, |
| "max_span_length": self.max_span_length, |
| } |
| for name, value in positive.items(): |
| if value <= 0: |
| raise ValueError(f"{name} must be positive, got {value}.") |
| for name in ( |
| "encoder_num_hidden_layers", |
| "pool_num_hidden_layers", |
| "latent_encoder_num_hidden_layers", |
| "latent_decoder_num_hidden_layers", |
| "dit_num_hidden_layers", |
| "decoder_num_hidden_layers", |
| "decoder_condition_num_hidden_layers", |
| ): |
| if getattr(self, name) < 0: |
| raise ValueError(f"{name} cannot be negative.") |
| if self.encoder_hidden_size % self.encoder_num_attention_heads: |
| raise ValueError( |
| "encoder_hidden_size must be divisible by " |
| "encoder_num_attention_heads." |
| ) |
| if ( |
| self.dit_num_attention_heads * self.dit_head_dim |
| != self.dit_hidden_size |
| ): |
| raise ValueError( |
| "dit_num_attention_heads * dit_head_dim must equal " |
| "dit_hidden_size." |
| ) |
| if ( |
| self.decoder_num_attention_heads * self.decoder_head_dim |
| <= 0 |
| ): |
| raise ValueError("Invalid decoder attention geometry.") |
| if ( |
| self.decoder_num_attention_heads |
| % self.decoder_num_key_value_heads |
| ): |
| raise ValueError( |
| "decoder_num_attention_heads must be divisible by " |
| "decoder_num_key_value_heads." |
| ) |
| if self.decoder_hidden_size % self.condition_num_attention_heads: |
| raise ValueError( |
| "decoder_hidden_size must be divisible by " |
| "condition_num_attention_heads." |
| ) |
| if not 0 < self.dit_rope_dim <= self.dit_head_dim: |
| raise ValueError("dit_rope_dim must be in (0, dit_head_dim].") |
| if self.dit_rope_dim % 2: |
| raise ValueError("dit_rope_dim must be even.") |
| if self.decoder_head_dim % 2: |
| raise ValueError("decoder_head_dim must be even.") |
| if not 1 <= self.min_span_length <= self.max_span_length: |
| raise ValueError( |
| "Span lengths must satisfy 1 <= min_span_length <= " |
| "max_span_length." |
| ) |
| if not 0 < self.target_compression_rate <= 1: |
| raise ValueError("target_compression_rate must be in (0, 1].") |
| if self.scaling_factor == 0: |
| raise ValueError("scaling_factor cannot be zero.") |
| if self.diffusion_time <= 0: |
| raise ValueError("diffusion_time must be positive.") |
|
|
| @classmethod |
| def tiny(cls, **overrides: Any) -> "Cola2Config": |
| """Return a CPU-friendly config that uses the production code path.""" |
|
|
| values: dict[str, Any] = { |
| "vocab_size": 67, |
| "encoder_embedding_size": 32, |
| "encoder_hidden_size": 32, |
| "encoder_intermediate_size": 64, |
| "encoder_num_attention_heads": 4, |
| "encoder_num_hidden_layers": 0, |
| "pool_num_hidden_layers": 1, |
| "latent_encoder_num_hidden_layers": 1, |
| "latent_decoder_num_hidden_layers": 0, |
| "latent_dim": 8, |
| "min_span_length": 1, |
| "max_span_length": 3, |
| "latent_block_size": 4, |
| "dit_hidden_size": 32, |
| "dit_intermediate_size": 64, |
| "dit_num_hidden_layers": 2, |
| "dit_num_attention_heads": 4, |
| "dit_head_dim": 8, |
| "dit_rope_dim": 8, |
| "dit_block_size": 4, |
| "decoder_hidden_size": 32, |
| "decoder_intermediate_size": 64, |
| "decoder_num_hidden_layers": 2, |
| "decoder_num_attention_heads": 4, |
| "decoder_num_key_value_heads": 2, |
| "decoder_head_dim": 8, |
| "condition_num_attention_heads": 4, |
| "decoder_condition_num_hidden_layers": 1, |
| "duration_hidden_size": 32, |
| "local_position_dim": 8, |
| "max_position_embeddings": 256, |
| "pad_token_id": 0, |
| "mask_token_id": 0, |
| "bos_token_id": 1, |
| "eos_token_id": 2, |
| } |
| values.update(overrides) |
| return cls(**values) |
|
|
|
|
| class Cola2GenerationConfig(GenerationConfig): |
| """Generation controls for the Full-DiT + span-causal inference path.""" |
|
|
| model_type = "cola2" |
|
|
| def __init__( |
| self, |
| dit_num_steps: int = 8, |
| dit_timestep_schedule: str = "low_noise", |
| dit_timestep_power: Optional[float] = None, |
| dit_cfg_disable_empty_prefix: bool = False, |
| dit_guidance_scale: float = 10.0, |
| decoder_guidance_scale: float = -0.9, |
| span_causal_chunk_size: int = 1, |
| dynamic_duration: bool = True, |
| latent_feedback: bool = True, |
| **kwargs: Any, |
| ) -> None: |
| |
| dit_num_steps = int(kwargs.pop("dit_steps", dit_num_steps)) |
| dit_guidance_scale = float( |
| kwargs.pop("dit_cfg_scale", dit_guidance_scale) |
| ) |
| decoder_guidance_scale = float( |
| kwargs.pop("decoder_cfg_scale", decoder_guidance_scale) |
| ) |
| span_causal_chunk_size = int( |
| kwargs.pop( |
| "span_causal_chunk_spans", |
| kwargs.pop("chunk_spans", span_causal_chunk_size), |
| ) |
| ) |
| latent_feedback = bool( |
| kwargs.pop("latent_feedback_reencode", latent_feedback) |
| ) |
| kwargs.setdefault("max_new_tokens", 32) |
| kwargs.setdefault("do_sample", False) |
| kwargs.setdefault("temperature", 1.0) |
| kwargs.setdefault("top_k", 50) |
| kwargs.setdefault("top_p", 1.0) |
| super().__init__(**kwargs) |
| self.dit_num_steps = dit_num_steps |
| self.dit_timestep_schedule = str(dit_timestep_schedule) |
| self.dit_timestep_power = ( |
| None |
| if dit_timestep_power is None |
| else float(dit_timestep_power) |
| ) |
| self.dit_cfg_disable_empty_prefix = bool( |
| dit_cfg_disable_empty_prefix |
| ) |
| self.dit_guidance_scale = dit_guidance_scale |
| self.decoder_guidance_scale = decoder_guidance_scale |
| self.span_causal_chunk_size = span_causal_chunk_size |
| self.dynamic_duration = bool(dynamic_duration) |
| self.latent_feedback = bool(latent_feedback) |
| self.validate() |
|
|
| @property |
| def dit_steps(self) -> int: |
| return self.dit_num_steps |
|
|
| @dit_steps.setter |
| def dit_steps(self, value: int) -> None: |
| self.dit_num_steps = int(value) |
|
|
| @property |
| def dit_cfg_scale(self) -> float: |
| return self.dit_guidance_scale |
|
|
| @dit_cfg_scale.setter |
| def dit_cfg_scale(self, value: float) -> None: |
| self.dit_guidance_scale = float(value) |
|
|
| @property |
| def decoder_cfg_scale(self) -> float: |
| return self.decoder_guidance_scale |
|
|
| @decoder_cfg_scale.setter |
| def decoder_cfg_scale(self, value: float) -> None: |
| self.decoder_guidance_scale = float(value) |
|
|
| @property |
| def chunk_spans(self) -> int: |
| return self.span_causal_chunk_size |
|
|
| @chunk_spans.setter |
| def chunk_spans(self, value: int) -> None: |
| self.span_causal_chunk_size = int(value) |
|
|
| @property |
| def span_causal_chunk_spans(self) -> int: |
| return self.span_causal_chunk_size |
|
|
| @span_causal_chunk_spans.setter |
| def span_causal_chunk_spans(self, value: int) -> None: |
| self.span_causal_chunk_size = int(value) |
|
|
| @property |
| def latent_feedback_reencode(self) -> bool: |
| return self.latent_feedback |
|
|
| @latent_feedback_reencode.setter |
| def latent_feedback_reencode(self, value: bool) -> None: |
| self.latent_feedback = bool(value) |
|
|
| def validate(self, is_init: bool = False) -> None: |
| super().validate(is_init=is_init) |
| if getattr(self, "dit_num_steps", 1) < 1: |
| raise ValueError("dit_num_steps must be at least 1.") |
| if getattr(self, "dit_timestep_schedule", "uniform") not in { |
| "uniform", |
| "high_noise", |
| "low_noise", |
| "power", |
| }: |
| raise ValueError( |
| "dit_timestep_schedule must be uniform, high_noise, " |
| "low_noise, or power." |
| ) |
| if ( |
| getattr(self, "dit_timestep_schedule", "uniform") == "power" |
| and ( |
| getattr(self, "dit_timestep_power", None) is None |
| or float(self.dit_timestep_power) <= 0 |
| ) |
| ): |
| raise ValueError( |
| "dit_timestep_power must be positive for power schedule." |
| ) |
| if getattr(self, "span_causal_chunk_size", 1) < 1: |
| raise ValueError("span_causal_chunk_size must be at least 1.") |
|
|
